PyTorch training loop cheat sheet

This sheet is for people who have written a PyTorch loop before and want the order of operations and the sharp edges on one page. Every block runs on a CPU in a few seconds with synthetic data. For the ideas behind the loop see /learn/beginner/04-training-loop-and-loss/ and /learn/beginner/07-your-first-neural-network/.

One epoch. The inner loop runs per batch. The scheduler and the validation pass run once per epoch, after the last optimiser step.

The canonical loop

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

torch.manual_seed(0)

# Synthetic regression data: y = 3 x0 - 2 x1 + noise
x = torch.randn(512, 2)
y = (3 * x[:, 0] - 2 * x[:, 1] + 0.1 * torch.randn(512)).unsqueeze(1)
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(2, 16), nn.ReLU(), nn.Linear(16, 1))
loss_fn = nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

for epoch in range(20):
    model.train()                              # dropout and batch norm in training mode
    for xb, yb in train_loader:
        optimizer.zero_grad(set_to_none=True)  # clear last step's gradients
        pred = model(xb)                       # forward pass
        loss = loss_fn(pred, yb)               # scalar loss
        loss.backward()                        # gradients land in .grad
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()                       # update weights
    scheduler.step()                           # once per epoch, after optimizer.step()

    model.eval()                               # inference mode for layers that care
    total, n = 0.0, 0
    with torch.no_grad():                      # no graph, no .grad, less memory
        for xb, yb in val_loader:
            total += loss_fn(model(xb), yb).item() * len(xb)
            n += len(xb)
    if epoch % 5 == 4:
        print(f"epoch {epoch + 1:2d}  val_mse {total / n:.4f}  lr {scheduler.get_last_lr()[0]:.5f}")

epoch  5  val_mse 0.4046  lr 0.00854
epoch 10  val_mse 0.0395  lr 0.00500
epoch 15  val_mse 0.0286  lr 0.00146
epoch 20  val_mse 0.0270  lr 0.00000

Line by line, the things that matter:

Line	Why it is there	What goes wrong without it
`model.train()`	Switches dropout on and batch norm to batch statistics	Validation leaks into running stats, or dropout stays off and you overfit
`optimizer.zero_grad(set_to_none=True)`	Gradients accumulate across `backward()` calls by design	Every step uses the sum of all previous gradients. Loss explodes or stalls
`loss.backward()`	Walks the autograd graph, fills `.grad` on every leaf parameter	Nothing to step with. Calling it twice on one graph raises an error unless `retain_graph=True`
`clip_grad_norm_`	Rescales all gradients so their global norm is at most `max_norm`	Rare large gradients from a bad batch wreck the weights. Clip after `backward()`, before `step()`
`optimizer.step()`	Applies the update rule to every parameter with a `.grad`	Weights never change
`scheduler.step()`	Moves the learning rate along its schedule	Calling it before `optimizer.step()` skips the first value and emits a warning. Per-step schedulers (`OneCycleLR`) go inside the batch loop
`model.eval()`	Dropout off, batch norm uses running stats	Validation numbers are noisy and too pessimistic
`torch.no_grad()`	Stops autograd recording	Validation builds a graph it never frees, memory climbs, and it runs slower
`.item()`	Pulls a Python float off the device	Keeping the tensor keeps its graph alive. Summing loss tensors across an epoch is a classic leak

torch.inference_mode() is a stricter and slightly faster no_grad(). Tensors created inside it cannot be used in autograd later, which is fine for evaluation and wrong for anything you will backpropagate through.

Mixed precision

Autocast runs matmuls and convolutions in a lower precision and keeps reductions, softmax and the master weights in fp32. On CPU the supported low precision is bf16, so this block runs anywhere. On CUDA, swap device_type and keep bf16 on Ampere or newer. If you must use fp16 (older GPUs), add a torch.amp.GradScaler and call scaler.scale(loss).backward(), scaler.step(optimizer), scaler.update() in place of the plain calls, and unscale before clipping with scaler.unscale_(optimizer).

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 10))
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.randn(32, 64)
target = torch.randint(0, 10, (32,))

# bf16 autocast works on CPU, so the same code runs on a laptop and a GPU.
# On CUDA use device_type="cuda"; for float16 add a torch.amp.GradScaler.
device_type = "cpu"
for step in range(3):
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
        logits = model(x)                       # matmuls run in bf16
        loss = nn.functional.cross_entropy(logits, target)  # reductions stay fp32
    loss.backward()                             # outside autocast
    optimizer.step()
    print(f"step {step}  loss {loss.item():.4f}  logits dtype {logits.dtype}  weights {model[0].weight.dtype}")

The output shows logits dtype torch.bfloat16 and weights torch.float32: activations are cast, parameters are not. backward() and step() sit outside the context on purpose.

Checkpoints

Save everything needed to resume, not only the weights. Load with weights_only=True (the default since PyTorch 2.6) so a checkpoint from an untrusted source cannot run code.

import tempfile, pathlib
import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(4, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)

# one fake step so the optimiser has state worth saving
loss = model(torch.randn(8, 4)).pow(2).mean()
loss.backward()
optimizer.step()
scheduler.step()

ckpt = {
    "epoch": 1,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "rng": torch.get_rng_state(),
}
path = pathlib.Path(tempfile.mkdtemp()) / "ckpt.pt"
torch.save(ckpt, path)

restored = nn.Linear(4, 2)
state = torch.load(path, map_location="cpu", weights_only=True)
restored.load_state_dict(state["model"])
same = all(torch.equal(a, b) for a, b in zip(model.state_dict().values(), restored.state_dict().values()))
print("epoch", state["epoch"], "weights identical:", same, "optimizer keys:", list(state["optimizer"].keys()))

Prints epoch 1 weights identical: True optimizer keys: ['state', 'param_groups']. For a model wrapped in DataParallel or DistributedDataParallel, save model.module.state_dict() so the keys do not carry a module. prefix. For sharing weights rather than resuming, prefer safetensors (see /reference/frameworks/).

Reproducibility

import random
import numpy as np
import torch
from torch import nn

def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)            # also seeds CUDA and MPS when present
    torch.use_deterministic_algorithms(True, warn_only=True)

def run(seed: int) -> float:
    seed_everything(seed)
    g = torch.Generator().manual_seed(seed)   # a generator for the DataLoader shuffle
    x = torch.randn(64, 8, generator=g)
    model = nn.Linear(8, 1)
    return model(x).sum().item()

print("same seed, same result:", run(1) == run(1))
print("different seed, different result:", run(1) != run(2))

Both lines print True. On CUDA you also need CUBLAS_WORKSPACE_CONFIG=:4096:8 in the environment for deterministic matmuls, and torch.backends.cudnn.benchmark = False. Bit-exact runs across GPU models or driver versions are not a reasonable target. Same machine, same software, same seed is.

DataLoader pitfalls

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

torch.manual_seed(0)
ds = TensorDataset(torch.arange(10).float().unsqueeze(1), torch.arange(10))

# drop_last=False leaves a short final batch, which skews a mean of per-batch losses
sizes = [len(xb) for xb, _ in DataLoader(ds, batch_size=4)]
print("batch sizes without drop_last:", sizes)

# a generator makes the shuffle order reproducible without touching the global seed
g = torch.Generator().manual_seed(123)
order = [int(yb[0]) for _, yb in DataLoader(ds, batch_size=1, shuffle=True, generator=g)]
print("shuffled order:", order)

# num_workers=0 on CPU demos; raise it only when the dataset does real decoding work
loader = DataLoader(ds, batch_size=5, num_workers=0, pin_memory=False)
print("batches per epoch:", len(loader))

batch sizes without drop_last: [4, 4, 2]
shuffled order: [2, 1, 6, 0, 4, 7, 9, 8, 3, 5]
batches per epoch: 2

Average losses by sample count, not by batch count, or set drop_last=True for training.
num_workers above 0 forks processes. On macOS and Windows the dataset and collate function must be importable (no lambdas), and the training script needs an if __name__ == "__main__": guard.
Each worker gets its own numpy state. Seed it in worker_init_fn or random augmentations repeat across workers.
pin_memory=True only helps when copying to CUDA. persistent_workers=True avoids re-spawning workers each epoch.
shuffle=True on the validation loader is harmless but pointless. shuffle=False on the training loader is a bug.

Common errors

Error or symptom	Cause	Fix
`RuntimeError: Expected all tensors to be on the same device`	Model on GPU, batch on CPU (or the reverse)	Move both: `model.to(device)` once, `xb.to(device)` per batch
`Trying to backward through the graph a second time`	A tensor from a previous step is reused in the loss (often a hidden state kept across batches)	`.detach()` what you carry over
Loss is `nan` after a few steps	Learning rate too high, fp16 overflow, log of zero	Lower the rate, use bf16 or a GradScaler, add epsilon inside logs
Loss stuck at `ln(num_classes)`	Labels are wrong dtype or shape, or `zero_grad` is missing	Labels as `int64` class ids of shape `(N,)`, logits `(N, C)`
Training loss falls, validation loss falls, test is bad	Leakage in the split	See /learn/beginner/02-data-is-the-model/
`CUDA out of memory` at validation	Validation outside `no_grad()`, or batch larger than training	Wrap in `torch.no_grad()`, match batch sizes
Memory grows every epoch	Accumulating loss tensors instead of `.item()` floats	Call `.item()` or `.detach()`
`UserWarning: Detected call of lr_scheduler.step() before optimizer.step()`	Order swapped, or the first step was skipped by a GradScaler	Step the optimiser first. The warning is harmless when the scaler skips a step
Results differ run to run	No seed, non-deterministic ops, worker seeding	The reproducibility block above
`view size is not compatible`	Tensor is non-contiguous after `permute`	`.contiguous()` or `.reshape()`, see /reference/tensor-shapes-and-broadcasting/