The canonical training loop

Five lines, in this order

Almost every PyTorch training loop has an outer loop over epochs and an inner loop over batches. Inside, the order is fixed:

1. optimizer.zero_grad() clears old gradients, because they accumulate.
2. logits = model(x) runs the forward pass and builds the graph.
3. loss = loss_fn(logits, y) reduces everything to one scalar.
4. loss.backward() fills .grad for every parameter.
5. optimizer.step() updates the parameters using those gradients.

Before these, move the batch to the same device as the model. Call model.train() at the start of each epoch, because evaluation code elsewhere may have switched the model to eval mode.

For logging, accumulate loss.item() (a Python float), never the loss tensor, which would keep every batch's graph in memory. zero_grad() may also go right after step(); what matters is that it happens between one step() and the next backward(). Swapping backward() and step() means each update uses the previous batch's gradients, or none.

Tracing one iteration

Follow one batch through the loop. The state shows the model mode, whether a graph exists, what .grad holds and whether the weights have changed.

for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        logits = model(x)
        loss = loss_fn(logits, y)
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch} last batch loss {loss.item():.4f}")

1. Line 1: The outer loop counts epochs: full passes over the training data.
2. Line 2: Switch dropout and batch norm to training behaviour. Validation at the end of the previous epoch may have left the model in eval mode.
3. Line 3: The DataLoader yields one shuffled mini-batch of inputs and labels.
4. Line 4: Data must be on the same device as the model's parameters, or the forward pass raises a device mismatch error.
5. Line 5: Old gradients are cleared first, because backward adds to whatever is already in `.grad`.
6. Line 6: The forward pass computes logits and autograd records every operation.
7. Line 7: The loss reduces the batch to one scalar, the root of the graph.
8. Line 8: Backward walks the graph in reverse, fills `.grad` for every parameter, then frees the graph.
9. Line 9: Only now do the weights change, using this batch's gradients.
10. Line 10: After all batches, the epoch summary prints. `.item()` gives a plain float, so no graph is kept alive.

A complete, self-contained example

Synthetic 2-D points labelled by which side of a line they fall on. Full-batch training keeps it short; the structure is the same as with a DataLoader.

import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(256, 2)
y = (X[:, 0] + X[:, 1] > 0).long()        # label: above or below the line

model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

for epoch in range(50):
    optimizer.zero_grad()
    loss = loss_fn(model(X), y)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    acc = (model(X).argmax(dim=1) == y).float().mean().item()
print(f"final loss {loss.item():.3f}, train accuracy {acc:.2%}")

Check the very first loss

With random initial weights, a C-class classifier predicts roughly uniform probabilities, so the first cross-entropy loss should be close to ln(C): about 2.30 for 10 classes, 0.69 for 2. A first loss far from that points to a bug in labels, the loss choice or the final layer.

Quick check: Which order of calls inside the batch loop is correct?

• `backward()`, `zero_grad()`, forward and loss, `step()`
• `zero_grad()`, forward and loss, `backward()`, `step()`
• forward and loss, `step()`, `backward()`, `zero_grad()`
• `step()`, forward and loss, `zero_grad()`, `backward()`