Five moves, repeated: forward → loss → backward → step → zero_grad. Write them once by hand, with nothing but raw tensors and autograd, and watch a model actually learn. This exact shape is in every PyTorch repo you'll ever open.
🎯 Mission: read & write real PyTorchTraining is Lesson 2's single gradient step, put in a loop over your data. Each pass: predict, measure how wrong you are, ask autograd for the gradients, nudge the parameters downhill, reset. Do it a few hundred times and the parameters slide into the values that fit the data. That loop — forward · loss · backward · step · zero_grad — is the spine of all of PyTorch.
Last lesson you took one gradient step on a toy expression. A real model just does that over and over on real data until the loss stops dropping. We'll learn a straight line — the simplest possible "model" — so the loop is the only new thing.
Why a line? Because pred = w·x + b is exactly what a single neural-network layer computes (a linear layer), just with more numbers. Master the loop on a line and scaling up to a real net changes the model, not the loop.
We make 100 points along a true line y = 2x − 3 with a little noise, then ask the model to recover w=2, b=−3 starting from zero — so we can check that learning worked.
torch.manual_seed(0) X = torch.linspace(-1, 1, 100).unsqueeze(1) # shape (100, 1) y = 2.0 * X - 3.0 + 0.1 * torch.randn_like(X) # the noisy truth w = torch.zeros(1, 1, requires_grad=True) # the parameters we'll learn, b = torch.zeros(1, requires_grad=True) # watched by autograd
X is (100, 1) and w is (1, 1) so X @ w is (100, 1) — one prediction per data point. .unsqueeze(1) turns a flat (100,) vector into a (100, 1) column so the matmul lines up. Shape bookkeeping is most of the work in real models — keep printing .shape.Here is the canonical loop. Read the five labelled lines — you will recognize them in every training script, just dressed up with library calls later.
lr = 0.1 for epoch in range(201): pred = X @ w + b # 1. forward loss = ((pred - y) ** 2).mean() # 2. loss (mean squared error → a scalar) loss.backward() # 3. backward: fills w.grad, b.grad with torch.no_grad(): # 4. step downhill (don't record the update) w -= lr * w.grad b -= lr * b.grad w.grad.zero_() # 5. zero grads for the next pass b.grad.zero_()
.mean()?
.backward() needs a single number. (pred - y)² is 100 errors; .mean() collapses them to one scalar — the mean squared error (MSE), the standard loss for predicting numbers. That one scalar is what we differentiate.Source: PyTorch — Optimization loop tutorial.
Printing the loss and parameters every 40 epochs (real output from the script on this machine):
epoch | loss | w | b
0 | 10.34831 | 0.136 | -0.599
40 | 0.01535 | 1.888 | -2.996
80 | 0.01051 | 1.993 | -2.996
120 | 0.01049 | 1.999 | -2.996
200 | 0.01049 | 2.000 | -2.996
learned w = 2.000 (true 2.0), b = -2.996 (true -3.0)
The loss plummets in the first ~40 epochs, then flattens once the parameters have basically found the line. The residual 0.0105 is the noise we baked in — the model correctly refuses to fit randomness. That refusal-to-overfit-noise is a good thing.
~/projects/learn/public/courses/pytorch/practice/.venv/bin/python ~/projects/learn/public/courses/pytorch/practice/lesson03_training_loop.pyThe learning rate lr scales every step. It's the hyperparameter you'll fiddle with most, so build intuition now:
| learning rate | what happens | symptom |
|---|---|---|
| too small (e.g. 0.0001) | tiny steps; barely moves | loss drops painfully slowly |
| just right (here 0.1) | fast, stable descent into the valley | loss falls smoothly, then flattens |
| too big (e.g. 5) | overshoots the minimum, bounces out | loss explodes to nan / oscillates |
nan loss? Lower the learning rate first
A loss that blows up to nan or inf is the classic "learning rate too high" signature. Cutting lr by 10× is the standard first move.Everything you wrote here has an idiomatic library replacement — that's Lesson 4. Keep this mapping in mind; it's how you'll read real code back into this loop:
| You wrote (from scratch) | Real code uses |
|---|---|
w, b + X @ w + b | nn.Linear(1, 1) — holds the params, does the matmul |
((pred - y)**2).mean() | nn.MSELoss() |
w -= lr * w.grad (per param) | optimizer.step() |
w.grad.zero_() (per param) | optimizer.zero_grad() |
nn.* and torch.optim.* so the same loop scales from a line to a million-parameter network..mean() on the squared errors before .backward()?nan after a few epochs. Most likely fix?nn.Linear(1, 1) replace?~0.0105 instead of 0. Why is that correct?.mean()?(pred − target)². .mean() turns the per-point errors into the single scalar that .backward() requires.nan. First fix for nan: cut lr 10×.w·x + b is exactly what one linear layer computes. The training loop is identical for a line or a deep net — only the model in step 1 grows.Try: "plot the line before and after training," "what changes if I shuffle or batch the data?", or "show me a loss that diverges and why." Next we make this loop idiomatic with nn.Module and torch.optim. Say "go" for Lesson 4.