training_with_GPU.py

num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    for x_batch,y_batch in train_dataloader:
        if train_on_gpu:
            x_batch,y_batch = x_batch.cuda(), y_batch.cuda()
        optimizer.zero_grad()
        pred = model(x_batch)
        loss = loss_criterion(pred, y_batch)
        loss.backward()
        optimizer.step()
    model.eval()
    for x_batch,y_batch in valid_dataloader:
        if train_on_gpu:
            x_batch,y_batch = x_batch.cuda(), y_batch.cuda()
        pred = model(x_batch)
        val_loss = loss_criterion(pred, y_batch)