The canonical five-line PyTorch training loop runs live against eight fixed data points scattered near y equals two x plus one. Each line of the loop highlights as its phase executes: model(x) makes eight predictions, mse_loss measures the average squared miss, zero_grad clears old gradients, backward computes dL/dw and dL/db, and step nudges w and b against the gradient. An amber model line on the scatter plot starts flat at zero and bends toward a faint dashed best-possible fit at w about 2.09 and b about 1.01, while a sparkline tracks the loss falling from 27.5 toward its floor of about 0.08. One button plays a single epoch slowly, phase by phase; another trains continuously at about 24 epochs per second. A logarithmic learning-rate slider from 0.001 to 1.5 sets the stride: around 0.08 the fit converges in some seventy epochs, while strides above roughly 0.2 make the line swing wildly and the loss explode until the run is declared diverged.
Five lines, all of training
A real fit, trained by the real loop. The amber line is the model ŷ = w·x + b, starting flat at zero. Watch each of the five lines do its one job — predict, measure, clear, differentiate, step — and nudge the line toward the data.
train.py
for epoch in range(200):
pred = model(x)
loss = F.mse_loss(pred, y)
opt.zero_grad()
loss.backward()
opt.step()
lr (stride)0.08
loss over epochs
epoch0
w0.000
b0.000
loss27.541
untrained
Every model you’ve ever heard of — GPT included — is trained by exactly these five lines, just with more numbers inside. backward() is the tape-replay from the autograd widget; step() is the stride from the gradient chapter; the loss is chapter six’s average surprise. The loop is where the whole math book reports for work.