Table of Contents
The Minimal Training Loop Mental Model
A PyTorch training loop is the repeated process of showing your model batches of data, measuring how wrong the predictions are, and adjusting the model parameters to reduce that error. The structure is always the same, even as models and datasets change. You initialize the model, loss function, and optimizer once, then you iterate over epochs, and inside each epoch you iterate over batches. Each batch does a forward pass to compute predictions, a loss computation, a backward pass to compute gradients, and an optimizer step to update weights.
Two modes matter in this structure. Training mode enables behaviors such as dropout and makes batch normalization update its running statistics, validation mode disables those training behaviors. You explicitly switch with model.train() and model.eval() and you also usually disable gradient tracking during validation so it runs faster and uses less memory.
The core order inside a training step is fixed: zero gradients, forward pass, compute loss, backward pass, optimizer step. If you change this order, you often get silent bugs.
One Epoch, Two Phases: Train Then Validate
A clean beginner friendly loop treats each epoch as having a training phase and an optional validation phase. During training you want gradients and parameter updates. During validation you want only evaluation of the current model, no gradient computation, no optimizer steps, and typically no randomness from training only layers.
In practice, the skeleton is: set model.train(), loop over training DataLoader, do update steps. Then set model.eval(), loop over validation DataLoader under torch.no_grad(), compute validation loss and metrics.
Always call model.eval() for validation and inference, and wrap validation code in with torch.no_grad():. Forgetting either can make validation noisy, slower, and can even change model state in batch normalization layers.
The Batch Update Step
Within the inner loop over batches, you do the same sequence every time. First, move the batch tensors to the same device as the model. Then reset gradients because PyTorch accumulates gradients by default. Then run the model on inputs to produce predictions. Then compute the loss by comparing predictions to targets. Then call loss.backward() to populate gradients in each parameter tensor. Finally, call optimizer.step() to update the parameters.
This step corresponds to minimizing an objective of the form
$$
\min_{\theta} \; \frac{1}{N}\sum_{i=1}^{N}\mathcal{L}(f_\theta(x_i), y_i),
$$
where $\theta$ are model parameters, $f_\theta$ is the model, and $\mathcal{L}$ is the loss function.
Gradients accumulate across backward calls. If you forget optimizer.zero_grad() (or set_to_none=True), your updates will effectively use summed gradients from multiple batches and training will behave incorrectly.
A Canonical PyTorch Skeleton
The following template is a solid default for beginner projects. It assumes you already have train_loader and optionally val_loader, plus a model, loss, and optimizer defined earlier in the chapter sequence.
import torch
def train_one_epoch(model, train_loader, optimizer, loss_fn, device):
model.train()
running_loss = 0.0
for x, y in train_loader:
x = x.to(device)
y = y.to(device)
optimizer.zero_grad(set_to_none=True)
y_pred = model(x)
loss = loss_fn(y_pred, y)
loss.backward()
optimizer.step()
running_loss += loss.item() * x.size(0)
epoch_loss = running_loss / len(train_loader.dataset)
return epoch_loss
@torch.no_grad()
def validate(model, val_loader, loss_fn, device):
model.eval()
running_loss = 0.0
for x, y in val_loader:
x = x.to(device)
y = y.to(device)
y_pred = model(x)
loss = loss_fn(y_pred, y)
running_loss += loss.item() * x.size(0)
epoch_loss = running_loss / len(val_loader.dataset)
return epoch_loss
def fit(model, train_loader, val_loader, optimizer, loss_fn, device, epochs):
model.to(device)
for epoch in range(epochs):
train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn, device)
if val_loader is not None:
val_loss = validate(model, val_loader, loss_fn, device)
print(f"epoch {epoch+1}/{epochs} train_loss={train_loss:.4f} val_loss={val_loss:.4f}")
else:
print(f"epoch {epoch+1}/{epochs} train_loss={train_loss:.4f}")
This structure makes two important ideas explicit. First, the model is moved to the device once in fit, not repeatedly. Second, training and validation are separated so you do not accidentally update weights during validation.
Where People Commonly Break the Structure
Beginners often compute loss correctly but forget that only optimizer.step() changes the model. Without it, the model never learns. Another common mistake is calling loss.backward() during validation or leaving model.train() on, which makes validation results inconsistent.
You should also be careful about what you average. loss.item() is a per batch scalar that may already be a mean over the batch, depending on the loss function reduction. Multiplying by batch_size and dividing by the dataset size is a reliable way to get an epoch average that is not biased by the final smaller batch.
Do not call optimizer.step() in validation, and do not compute validation loss with gradients enabled. Validation should never modify model parameters or accumulate gradients.
A Quick Checklist for a Healthy Loop
At the start of each training epoch you set training mode. For every batch you move data to the device, clear gradients, compute predictions, compute loss, backpropagate, then step the optimizer. After training, you switch to eval mode for validation, disable gradients, and only compute losses and metrics. If you can recite that sequence and your code matches it, you have the correct training loop structure that scales to nearly every model you will build later in the course.