def train(num_epoch):
    train_losses = []
    test_losses = []
    for i in range(num_epoch):
        for X_batch, y_batch in batch_generator(X_train, y_train, 500):
            # На лекции мы рассказывали, что дропаут работает по-разному во время обучения и реального предсказания
            # Чтобы это учесть нам нужно включать и выключать режим обучения, делается это командой ниже
            model.train(True)

            # Посчитаем предсказание и лосс
            y_pred = model(torch.from_numpy(X_batch))
            loss = loss_fn(y_pred, torch.from_numpy(y_batch))

            # зануляем градиент
            optimizer.zero_grad()

            # backward
            loss.backward()

            # ОБНОВЛЯЕМ веса
            optimizer.step()

            # Запишем число (не тензор) в наши батчевые лоссы
            train_losses.append(loss.item())

            # Теперь посчитаем лосс на тесте
            model.train(False)
            # Сюда опять же надо положить именно число равное лоссу на всем тест датасете
            Y_pred = model(X_train)
            test_losses.append(loss_fn(Y_pred,y_train).item())

    return train_losses, test_losses