The PyTorch training loop

Every PyTorch training script is a version of one loop. It feeds a batch through the model, measures the error, computes gradients and nudges the weights. Then it checks how well the model does on data it has not trained on. Once you can write this loop from memory, most framework questions become easy.

The full loop

This is a complete, runnable script. It trains a small network to tell whether 20 random features sum to a positive number:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)
x = torch.randn(512, 20)
y = (x.sum(dim=1) > 0).long()    # label is 1 when the 20 features sum to a positive number

train_loader = DataLoader(TensorDataset(x[:400], y[:400]), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(x[400:], y[400:]), batch_size=64)

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 2))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    for xb, yb in train_loader:
        opt.zero_grad(set_to_none=True)
        loss = loss_fn(model(xb), yb)
        loss.backward()
        opt.step()

    model.eval()
    correct = total = 0
    with torch.inference_mode():
        for xb, yb in val_loader:
            correct += (model(xb).argmax(dim=1) == yb).sum().item()
            total += yb.size(0)
    print(f"epoch {epoch}  val accuracy {correct / total:.3f}")

Output:

epoch 0  val accuracy 0.580
epoch 1  val accuracy 0.679
epoch 2  val accuracy 0.759
epoch 3  val accuracy 0.795
epoch 4  val accuracy 0.830
epoch 5  val accuracy 0.875
epoch 6  val accuracy 0.893
epoch 7  val accuracy 0.893
epoch 8  val accuracy 0.893
epoch 9  val accuracy 0.884

The accuracy climbs and then levels off near 0.89. The exact numbers can shift a little across PyTorch versions.

Why the order matters

Each step of the inner loop does one job, and the order is fixed:

1. zero_grad clears the gradients left from the last step.
2. The forward pass produces predictions.
3. The loss compares predictions with labels.
4. backward fills .grad for every parameter.
5. step moves each weight against its gradient.

Swap two of these and the run still produces numbers, but they are wrong. Forgetting step 1 stacks old gradients on new ones.

Train mode and eval mode

model.train() and model.eval() change how some layers behave. Dropout is active in train mode and off in eval mode. BatchNorm uses batch statistics in train mode and stored running statistics in eval mode.

Validation needs both switches: model.eval() for the layers, and torch.inference_mode() so no gradients are tracked. Together they make predictions deterministic and cheap.

When a run does nothing

Check three things in order.

1. Shapes. Print each tensor's shape for one batch. A silent broadcast can train the wrong function.
2. Gradient flow. After loss.backward(), look for missing or zero gradients:

for n, p in model.named_parameters():
    if p.grad is None:
        print("NO GRAD", n)
    elif p.grad.abs().sum() == 0:
        print("ZERO GRAD", n)

On this model the check prints nothing, because every weight got a gradient. A forgotten detach() or a wrapping no_grad() shows up here as a NO GRAD line.

3. Batch-size scale. CrossEntropyLoss uses reduction="mean" by default, so the loss does not grow with batch size. With reduction="sum", doubling the batch doubles the loss and gradient, which is the same as doubling the learning rate.

Common mistakes

• Forgetting model.eval(). Dropout stays on during validation, and the scores come out noisy.
• Calculating metrics with gradients on. It wastes memory. Use inference_mode().
• Changing batch size with reduction="sum" and the same learning rate. You have changed the step size without meaning to.

Check yourself

1. What goes wrong if you skip zero_grad? Answer: each step adds the new gradient to the old ones, so the updates are wrong.
2. Why must you call model.eval() before validation? Answer: it turns off dropout and makes BatchNorm use its running statistics, so predictions are stable.

Key takeaways

• The loop is zero_grad, forward, loss, backward, step. Keep that order.
• Validate in eval mode under inference_mode().
• When training stalls, check shapes first, then gradient flow, then the loss scale.
• Overfit one batch before trusting a full run.

Further learning

Next, speed the loop up and make it resumable in mixed precision and checkpoints.