Katkat
Back

Implement PyTorch training loop

AmazonApplied Scientist · Onsite
MediumApplied ScientistOnsiteCoding & Algorithms

Last updated

This question evaluates practical implementation skills in PyTorch, focusing on model and device management, batch-wise training mechanics, gradient handling, and optimizer interaction.

Implement a basic PyTorch training loop

You are given a PyTorch neural network model, a DataLoader that yields (inputs, targets) batches, an optimizer, and a loss function.

Write a function train(model, train_loader, optimizer, loss_fn, device, num_epochs) that:

  • Moves the model and input batches to the specified device (CPU or GPU).
  • Runs for num_epochs epochs.
  • For each batch, performs a forward pass, computes the loss, runs backpropagation, and updates the model parameters.
  • Properly zeros gradients at the right time.
  • Optionally prints or returns the average training loss per epoch.

Clearly show the order of operations inside the training loop (zeroing gradients, forward pass, loss computation, backward pass, optimizer step).