Own the state, own the loop

Forward, loss, zero the gradients, backward, step. Written out, a PyTorch training loop is barely longer than its own description, and every line in it earns its place. The forward pass computes a prediction; the loss compares it to a target; then three more lines manage state that nothing in the math above required you to think about.

opt.zero_grad() exists because of a fact chapter two already put on the table: gradients accumulate into .grad across calls to backward(), they never clear themselves. Skip that line and the second step's gradient is the first step's gradient plus the second, and the third step's is worse than that. The loop has to clear the accumulator by hand, every time, because the tape has no opinion about when one step ends and the next one starts.

opt.step() is where the parameters actually move, and it makes that move under torch.no_grad() internally, mutating tensors in place without asking autograd to track the mutation. Do that write by hand instead, in place on a leaf tensor that still requires grad, and you've built exactly the kind of mistake this site's mistakes gallery keeps a specimen of. The optimizer's whole job is doing that unsafe-looking write safely, once, in the one place it belongs.

The loop is explicit because the state is yours.

None of this is theoretical. The loop below fits a small MLP to a synthetic regression target, two hundred steps, Adam at a fixed learning rate, and it runs to a specific loss on this machine: 0.00372, not an order-of-magnitude guess. Read the five lines inside the for loop against the previous section and each one should already have a job attached to it before you see what it prints.

Adam itself is state, though the loop above never shows it directly. Every parameter gets a running estimate of its gradient's mean and its squared magnitude, both updated on opt.step(), both invisible unless you go looking for them in opt.state_dict(). Nothing about that memory is optional. It's what makes Adam behave like Adam instead of like plain gradient descent wearing Adam's name.

import torch
from torch import nn

torch.manual_seed(0)
x = torch.randn(256, 8)
y = x.sum(dim=1, keepdim=True)

model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.Adam(model.parameters(), lr=1e-2)

for step in range(200):
    loss = nn.functional.mse_loss(model(x), y)
    opt.zero_grad()
    loss.backward()
    opt.step()
print(f"final loss {loss.item():.5f}")   # final loss 0.00372

Save model.state_dict() alone and reload it later, and the weights come back correct. What doesn't come back is Adam's memory: the running mean and variance it had built up for every parameter, gone, replaced by whatever a fresh opt.init would produce. Training resumes, technically, but from weights that used to belong to a well-adapted optimizer and now belong to a naive one. The loss curve after resume will not match the curve that never stopped.

The fix is to save the whole state as one unit: model.state_dict(), opt.state_dict(), and the step count, together, in one dictionary, one file. If bitwise resume matters, meaning the run after restore should reproduce identical numbers rather than merely plausible ones, the RNG state belongs in that same dictionary too. This is the exact rule the jax path's twelfth chapter states for its own training loop; only the serialization surface differs, state_dict and torch.save here instead of a pytree and Orbax there.

Nothing about torch.load or load_state_dict is unusual once the save side is right. The real discipline was deciding, before the first checkpoint ever got written, what counts as the state in the first place.

import torch

torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, "ckpt.pt")
state = torch.load("ckpt.pt")
model.load_state_dict(state["model"])
opt.load_state_dict(state["opt"])

A learning-rate schedule is not a separate mechanism bolted onto training. torch.optim.lr_scheduler objects wrap an existing optimizer, and calling scheduler.step() once per epoch or once per step mutates that optimizer's lr in place: same optimizer, same state dict, one more field changing over time.

Nothing here contradicts the chapter's opening claim. A scheduler is state you own, wired through the same loop you already wrote, mutated by an explicit call in exactly the shape of zero_grad and step. PyTorch never hides this behind a config file; it hands you the object and expects you to call it yourself.