The training loop

A complete, correct training and evaluation loop — loss, optimiser, batching, and validation that you can trust.

The canonical loop

import torch
from torch import nn

model = MyNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for epoch in range(EPOCHS):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)

        optimizer.zero_grad()          # 1. clear old gradients
        logits = model(xb)             # 2. forward
        loss = criterion(logits, yb)   # 3. measure error
        loss.backward()                # 4. backpropagate
        optimizer.step()               # 5. update parameters

    model.eval()                       # switch off dropout / batchnorm training mode
    correct = total = 0
    with torch.no_grad():
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            preds = model(xb).argmax(dim=1)
            correct += (preds == yb).sum().item()
            total += yb.size(0)
    print(f"epoch {epoch}: val acc {correct / total:.3f}")

💡

The five-line step order is not stylistic. Swapping any two — most commonly forgetting zero_grad() or calling step() before backward() — produces a model that silently trains badly rather than an error.

Making it converge

Symptom	Likely cause	Fix
Loss = NaN	Learning rate too high, bad normalisation	Lower LR, clip gradients, check inputs
Loss flat	LR far too low, dead activations	Raise LR, check layer init
Train great, val poor	Overfitting	More data, dropout, weight decay, early stop
Validation loss rising	Overfitting mid-training	Early stopping on val loss
Nothing learns	Labels or inputs mismatched	Verify shapes and that labels align

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)
# ... at the end of each epoch:
scheduler.step()

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)   # stabilise

Start with a small learning rate schedule you understand and a modest model. Most "the architecture is wrong" problems are actually an unstable or mis-set learning rate.

Reproducibility and honest metrics

import random, numpy as np

def seed_everything(seed=42):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

Seed everything, and still report results over several runs — deep learning is stochastic.
Track the validation metric every epoch and keep the best checkpoint, not the last one.
Never tune on the test set; use validation for decisions and the test set exactly once.
Save the optimiser state alongside the model if you intend to resume training.

FAQ

Adam or SGD?

Adam/AdamW converge quickly with little tuning and are the sensible default. SGD with momentum can generalise slightly better on some vision tasks once tuned — a research decision, not a starting point.

How large should the batch be?

As large as fits in memory, then adjust the learning rate with it. Larger batches are more efficient but often need a higher LR or warmup.

Modules, data and saving Evaluation and overfitting

Last refreshed 2026-09-18.

The training loop

The canonical loop

Making it converge

Reproducibility and honest metrics

FAQ

Related