I am trying to do batch learning on a large dataset that will not fit on the GPU.  I am not sure where to clear the gradients and compute the loss.  Is this the correct way to use a DataLoader and move the data to the GPU in pieces for batch learning?
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size, shuffle=True)

for epoch in range(n_epochs):
    for loc_X_train,loc_y_train in train_dataloader:
        # move data to the device
        loc_X_train, loc_y_train = loc_X_train.cuda(), loc_y_train.cuda()

        # clear any calculated gradients
        optimizer.zero_grad()

        # forward pass, compute outputs
        outputs = model.forward(loc_X_train)

        # compute loss
        loss = loss_function(outputs, loc_y_train)

        # backward pass, compute gradients
        loss.backward()

    # update learnable parameters
    optimizer.step()