у меня есть нейросеть и функция её обучения:

"

# свертка - релу - свертка - релу - свертка - релу - пулинг - батчнорм - полносвязный слой - релу - полносвязный слой

class CustomModel(nn.Module):

  def __init__(self):
    super(CustomModel, self).__init__()

    # свертка - релу - 1
    # размерность входа 3, тк rgb
    self.conv_1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, stride = 1, padding = 0)
    self.relu_1 = nn.ReLU()

    # свертка - релу - 2
    self.conv_2 = nn.Conv2d(in_channels = 32, out_channels = 64, kernel_size = 3, stride = 1, padding = 0)
    self.relu_2 = nn.ReLU()

    # свертка - релу - 3
    self.conv_3 = nn.Conv2d(in_channels = 64, out_channels = 128, kernel_size = 3, stride = 1, padding = 0)
    self.relu_3 = nn.ReLU()

    # пулинг
    self.pooling = nn.MaxPool2d(4)
    self.batch_norm = nn.BatchNorm2d(128)

    # полносвязный слой - 1
    self.linear_1 = nn.Linear(128 * 6 * 6, 256)
    # откуда берётся 6 * 6 - ответ:
    # 1) после первого слоя получили ((32 - 3 + 2 * 0) / 1) + 1 = 30
    # 2) после второго слоя получили ((30 - 3 + 2 * 0) / 1) + 1 = 28
    # 3) после третьего слоя получили ((28 - 3 + 2 * 0) / 1) + 1 = 26
    # 4) после пулинга получили 26 / 4 = 6 (в меньшую сторону округляется)
    # линейный слой принимает не тензор, а массив, поэтому (128, 6, 6) надо превартить в 1D размерность

    # релу
    self.relu_4 = nn.ReLU()

    # полносвязный слой - 2
    self.linear_2 = nn.Linear(256, 10)

  def forward(self, x):
    x = self.relu_1(self.conv_1(x))

    x = self.relu_2(self.conv_2(x))

    x = self.relu_3(self.conv_3(x))

    x = self.batch_norm(self.pooling(x))

    x = self.linear_1(x)
    x = self.relu_4(x)
    x = self.linear_2(x)

    return x

"

"
def train(
    model,
    criterion,
    optimizer,
    train_batch_gen,
    val_batch_gen,
    num_epochs=40
):
    '''
    Функция для обучения модели и вывода лосса и метрики во время обучения.

    :param model: обучаемая модель
    :param criterion: функция потерь
    :param optimizer: метод оптимизации
    :param train_batch_gen: генератор батчей для обучения
    :param val_batch_gen: генератор батчей для валидации
    :param num_epochs: количество эпох

    :return: обученная модель
    :return: (dict) accuracy и loss на обучении и валидации ("история" обучения)
    '''

    history = defaultdict(lambda: defaultdict(list))

    for epoch in tqdm(range(num_epochs)):
        train_loss = 0
        train_acc = 0
        val_loss = 0
        val_acc = 0

        start_time = time.time()

        # Устанавливаем поведение dropout / batch_norm  в обучение
        model.train(True)

        # На каждой "эпохе" делаем полный проход по данным
        for X_batch, y_batch in train_batch_gen:
            # Обучаемся на батче (одна "итерация" обучения нейросети)
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)

            # Логиты на выходе модели
            logits = model(X_batch)

            # Подсчитываем лосс
            loss = criterion(logits, y_batch.long().to(device))

            # Обратный проход
            loss.backward()
            # Шаг градиента
            optimizer.step()
            # Зануляем градиенты
            optimizer.zero_grad()

            # Сохраняем лоссы и точность на трейне
            train_loss += loss.detach().cpu().numpy()
            y_pred = logits.max(1)[1].detach().cpu().numpy()
            train_acc += np.mean(y_batch.cpu().numpy() == y_pred)

        # Подсчитываем лоссы и сохраняем в "историю"
        train_loss /= len(train_batch_gen)
        train_acc /= len(train_batch_gen)
        history['loss']['train'].append(train_loss)
        history['acc']['train'].append(train_acc)

        # Устанавливаем поведение dropout / batch_norm в режим тестирования
        model.train(False)

        # Полный проход по валидации
        for X_batch, y_batch in val_batch_gen:
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)

            # Логиты, полученные моделью
            logits = model(X_batch)

            # Лосс на валидации
            loss = criterion(logits, y_batch.long().to(device))

            # Сохраняем лоссы и точность на валидации
            val_loss += loss.detach().cpu().numpy()
            y_pred = logits.max(1)[1].detach().cpu().numpy()
            val_acc += np.mean(y_batch.cpu().numpy() == y_pred)

        # Подсчитываем лоссы и сохраняем в "историю"
        val_loss /= len(val_batch_gen)
        val_acc /= len(val_batch_gen)
        history['loss']['val'].append(val_loss)
        history['acc']['val'].append(val_acc)

        clear_output(wait=True)

        # Печатаем результаты после каждой эпохи
        print("Epoch {} of {} took {:.3f}s".format(
            epoch + 1, num_epochs, time.time() - start_time))
        print("  training loss (in-iteration): \t{:.6f}".format(train_loss))
        print("  validation loss (in-iteration): \t{:.6f}".format(val_loss))
        print("  training accuracy: \t\t\t{:.2f} %".format(train_acc * 100))
        print("  validation accuracy: \t\t\t{:.2f} %".format(val_acc * 100))

    return model, history
"

используя эти данные, сделай эти 2 пункта:

3. Посчитайте среднее время одного прямого прохода (forward pass) и среднее время обратного прохода (backward pass) по сети во время обучения.
4. Сохраните историю изменения ошибки (loss) и точности (accuracy) во время обучения.