PyTorch From First Principles
Chapter 15 of 16

What Are We Actually Doing?

By AI, PyTorch

Explain this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

Apply this chapter with AI

Copy this prompt into ChatGPT, Claude, Gemini, a local model, or another AI.

It has never been easier to obtain working PyTorch code. Describe a model to an assistant, paste the result into a file, and a training loop will usually appear and usually run. For a large fraction of everyday work this is fine, and this book does not ask you to pretend otherwise.

The difficulty arrives afterwards.

The code runs, but the loss sits at exactly the same value for two hundred steps. Or one tensor is [32, 128] where the next operation wanted [128, 32]. Or the gradient of a weight is None and nothing says why. Or the model trained yesterday and today it does not, with no change you can point to. None of these are syntax errors. They do not produce a helpful traceback, and often they produce no traceback at all. They are questions about what the program is doing, and to answer them you need somewhere to look.

That is the gap this book exists to close. The aim is not to memorize PyTorch’s API surface, and not to be able to reproduce a tutorial from memory. It is to understand the machinery well enough that unfamiliar PyTorch code becomes something you can interrogate: read it, form a hypothesis about what it should do, find the evidence that confirms or refutes the hypothesis, and repair the actual cause rather than the visible symptom.

You will use AI throughout that process, and you should. An assistant is excellent at producing a candidate explanation, a trace, a comparison, or an experiment. It is much weaker at knowing which of its explanations is true in your particular program. Deciding that remains your job, and it is a job you can only do if you know what the framework is actually doing.

Everything in the book is built on one loop, and this chapter is about that loop alone. Before any of PyTorch’s abstractions make it look complicated:

What are we actually doing when we train a model?

The smallest problem that is still training

Set neural networks aside. Suppose the world obeys a rule:

y = 3x

If x is 2, then y is 6. We are going to pretend we do not know the multiplier. All we know is the form of the rule:

y = wx

The value w is the one number the program is allowed to change. That makes it a parameter: not an input, not an output, but an internal value that training is permitted to adjust. Everything else in this section is fixed.

Start with a deliberately bad guess, w = 1. Feed in x = 2:

prediction = 2 × 1 = 2

The correct answer, the target, is 6. The model is wrong, and we now need to say how wrong in a way a computer can act on. A single number that gets larger the worse the prediction is will do. Squared error is the usual first choice:

loss = (prediction - target)²
     = (2 - 6)²
     = 16

Squaring does two useful things. It removes the sign, so being four too low and four too high are equally bad, and it grows faster than the error itself, so large mistakes dominate small ones. It also has a property that matters more than either: it is smooth, so we can differentiate it.

That matters because we now need a direction. We could try random values of w until the loss got smaller, and with a single parameter that would even work. It stops working the moment there are more than a handful: a small model has thousands of parameters and a large one has billions, and searching that space by trial is hopeless. What we want instead is a local question.

If I nudge w upward slightly, does the loss go up or down, and by how much?

That is the gradient: the derivative of the loss with respect to the parameter. Write the loss out in terms of w:

loss(w) = (wx - y)²

and differentiate it:

dloss/dw = 2(wx - y) · x

Substituting w = 1, x = 2, y = 6:

dloss/dw = 2(1·2 - 6) · 2
         = 2(-4) · 2
         = -16

The gradient is -16. Read that carefully, because the sign is the whole point. A negative gradient means that increasing w decreases the loss. So we should move w up. The magnitude, 16, tells us the slope is steep here, but it does not tell us how far to go; the derivative only describes the loss surface in the immediate neighborhood of the current value. To turn a direction into a step we introduce a scale factor, the learning rate:

w_new = w - learning_rate × gradient

With a learning rate of 0.1:

w_new = 1 - 0.1 × (-16)
      = 1 + 1.6
      = 2.6

The guess moved from 1 toward 3. Check what that buys us:

prediction = 2 × 2.6 = 5.2
loss       = (5.2 - 6)² = 0.64

The loss fell from 16 to 0.64 in one step. That is the entire algorithm. Predict, measure the error, compute the slope of the error with respect to each parameter, take a small step against the slope, repeat.

    flowchart LR
    A[input x] --> B[model: prediction = f_w of x]
    B --> C[loss: compare prediction with target y]
    C --> D[gradient: d loss / d w]
    D --> E[update: w = w - lr * gradient]
    E --> B

Everything the rest of this book adds is scale and organization. w becomes millions of numbers. x * w becomes matrix multiplications, activations, convolutions, attention. The hand-written update becomes an optimizer with momentum and per-parameter scaling. The shape of the loop does not change.

The same calculation, in PyTorch

Install PyTorch with the command for your operating system and hardware from the official site, then run this:

import torch

x = torch.tensor(2.0)
target = torch.tensor(6.0)
w = torch.tensor(1.0, requires_grad=True)

prediction = x * w
loss = (prediction - target) ** 2

print("prediction:", prediction.item())
print("loss:", loss.item())
print("w.grad before backward:", w.grad)

loss.backward()

print("w.grad after backward:", w.grad.item())

The output:

prediction: 2.0
loss: 16.0
w.grad before backward: None
w.grad after backward: -16.0

PyTorch computed -16.0. We computed -16 by hand a moment ago. That agreement is not a formality; it is the first piece of evidence in this book, and the habit it represents matters more than the result. Whenever PyTorch does something that looks like magic, there is a mechanism underneath that you can check against something you worked out yourself.

Three lines deserve individual attention.

w = torch.tensor(1.0, requires_grad=True). In the mathematical sense, w is our model parameter because it is a value that training is allowed to change. In PyTorch terms, it is currently a leaf tensor with requires_grad=True. Later, when we build models with nn.Module, we will meet nn.Parameter, a special kind of tensor that PyTorch registers as model state. For now, requires_grad=True means: track operations involving this tensor because I will want derivatives with respect to it.

x and target are also tensors, but they carry requires_grad=False, which is the default.

prediction = x * w. Because one operand requires gradients, PyTorch does not just compute the number and discard everything else. It also records that this particular result came from multiplying these particular tensors. The same happens for the subtraction and the squaring. By the time loss exists, PyTorch holds a record of the chain of operations that produced it.

loss.backward(). This walks that record backwards, applying the chain rule at each step. By default, the resulting gradients accumulate into the .grad attributes of leaf tensors that require gradients, such as w. PyTorch also computes the intermediate gradients needed to continue the backward pass, but it does not retain those intermediate .grad values unless you explicitly ask it to.

You can see the difference between the trainable leaf tensor w and a tensor produced by the computation directly:

print("w:      value=%s requires_grad=%s is_leaf=%s grad_fn=%s"
      % (w.item(), w.requires_grad, w.is_leaf, w.grad_fn))
print("x:      value=%s requires_grad=%s is_leaf=%s grad_fn=%s"
      % (x.item(), x.requires_grad, x.is_leaf, x.grad_fn))
print("loss:   value=%s requires_grad=%s is_leaf=%s grad_fn=%s"
      % (loss.item(), loss.requires_grad, loss.is_leaf, loss.grad_fn))

w:      value=1.0 requires_grad=True is_leaf=True grad_fn=None
x:      value=2.0 requires_grad=False is_leaf=True grad_fn=None
loss:   value=16.0 requires_grad=True is_leaf=False grad_fn=<PowBackward0 object at ...>

w requires gradients and was created directly rather than computed, which is what leaf means; it is an endpoint of the record, and it is exactly the kind of tensor whose .grad gets filled in. loss requires gradients too, but it was computed, so it carries a grad_fn, a reference to the operation that produced it. That grad_fn is the thread backward() pulls on. Chapter 3 takes this apart properly. For now, one distinction is enough to carry forward:

A tensor’s value and a tensor’s gradient are two different things, stored in two different places. w holds 1.0. w.grad holds -16.0. Confusing the two is the source of a surprising number of training bugs.

backward() computes gradients; it does not train anything

It is easy to read loss.backward() as the line that does the learning. It is not. Watch what happens to w:

x = torch.tensor(2.0)
target = torch.tensor(6.0)
w = torch.tensor(1.0, requires_grad=True)

before = w.item()
loss = (x * w - target) ** 2
loss.backward()

print("w before backward:", before)
print("w after backward: ", w.item())
print("w.grad:           ", w.grad.item())

w before backward: 1.0
w after backward:  1.0
w.grad:            -16.0

The parameter did not move. backward() filled in a derivative and stopped. Changing the parameter is a separate operation that you have not yet written.

So write it. The obvious attempt fails:

w -= 0.1 * w.grad

RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

The error is telling us something real. w is a leaf that requires gradients. If PyTorch recorded this subtraction the way it recorded the multiplication, w would become a computed tensor with a grad_fn of its own, and it would stop being the endpoint that gradients accumulate into. But this subtraction is not part of the model at all; it is us reaching in from outside and editing a parameter between iterations. Rather than quietly corrupting the structure, PyTorch refuses and asks us to state explicitly that this edit is not part of the computation:

learning_rate = 0.1

with torch.no_grad():
    w -= learning_rate * w.grad

print("w after update:", w.item())
print("still a parameter:", w.requires_grad, w.is_leaf)

new_prediction = x * w
print("new prediction:", new_prediction.item())
print("new loss:", ((new_prediction - target) ** 2).item())

w after update: 2.5999999046325684
still a parameter: True True
new prediction: 5.199999809265137
new loss: 0.6400002837181091

2.6, 5.2, 0.64. The same numbers we worked out on paper, allowing for float32 rounding. Inside torch.no_grad(), PyTorch performs the arithmetic without recording it, w is modified in place, and it remains the same parameter object it was before: still requiring gradients, still a leaf.

We now have three distinct phases, and it is worth naming them because the rest of the book keeps returning to this decomposition:

forward   prediction and loss are computed; the record is built
backward  the record is traversed; .grad is populated
update    parameters are modified; the record is not involved

Later you will replace the update with optimizer.step() and the manual gradient clearing with optimizer.zero_grad(). In current PyTorch, optimizer.zero_grad() defaults to setting gradients to None rather than filling them with zeros, so the exact state you observe before the next backward pass may differ from this hand-written loop. The three-phase structure does not. When a model refuses to learn, the useful first question is almost always which of these three phases is not doing its job, and you can only ask that question if you have kept them separate in your head.

Repeating the step

One step is not training. Put it in a loop:

import torch

x = torch.tensor(2.0)
target = torch.tensor(6.0)
w = torch.tensor(1.0, requires_grad=True)

learning_rate = 0.1

for step in range(8):
    prediction = x * w
    loss = (prediction - target) ** 2

    loss.backward()

    print(f"step={step} w={w.item():+.4f} loss={loss.item():9.4f} "
          f"w.grad={w.grad.item():+.4f}")

    with torch.no_grad():
        w -= learning_rate * w.grad

    w.grad.zero_()

step=0 w=+1.0000 loss=  16.0000 w.grad=-16.0000
step=1 w=+2.6000 loss=   0.6400 w.grad=-3.2000
step=2 w=+2.9200 loss=   0.0256 w.grad=-0.6400
step=3 w=+2.9840 loss=   0.0010 w.grad=-0.1280
step=4 w=+2.9968 loss=   0.0000 w.grad=-0.0256
step=5 w=+2.9994 loss=   0.0000 w.grad=-0.0051
step=6 w=+2.9999 loss=   0.0000 w.grad=-0.0010
step=7 w=+3.0000 loss=   0.0000 w.grad=-0.0002

Each row describes one coherent state, because the printing happens after backward() and before the update: the value of w used for this prediction, the loss that value produced, and the gradient computed from that loss. Three columns, three different things.

Read them together rather than separately. w climbs toward 3. The loss collapses toward zero. The gradient shrinks toward zero as well, which tells us we are arriving at a flat region of the loss surface; for this problem that flat region is the correct answer. Each column tells you something the others do not, and later in the book we will meet failures where the loss looks healthy and the gradients do not, or the reverse.

There is one line in that loop we have not justified.

Break it: remove w.grad.zero_()

Delete the last line of the loop body and run it again. Nothing raises an exception.

step=0 w=+1.0000 loss=  16.0000 w.grad=-16.0000
step=1 w=+2.6000 loss=   0.6400 w.grad=-19.2000
step=2 w=+4.5200 loss=   9.2416 w.grad=-7.0400
step=3 w=+5.2240 loss=  19.7847 w.grad=+10.7520
step=4 w=+4.1488 loss=   5.2790 w.grad=+19.9424
step=5 w=+2.1546 loss=   2.8591 w.grad=+13.1789
step=6 w=+0.8367 loss=  18.7199 w.grad=-4.1277
step=7 w=+1.2494 loss=  12.2578 w.grad=-18.1322

This is what a particularly difficult class of training failure looks like. There is no error message. The first step even works: the loss falls from 16 to 0.64, exactly as before. Then the run wanders, and after eight steps w is further from the answer than the initial guess was.

Resist the urge to change anything yet. The instinct at this point is to reach for the learning rate, because the run looks unstable and instability looks like a step-size problem. Instead, find a number you can predict and check it.

We know how to compute the gradient by hand. At step 1, w is 2.6:

dloss/dw = 2(wx - y) · x
         = 2(2.6 × 2 - 6) × 2
         = 2(-0.8) × 2
         = -3.2

The correct loop printed -3.2 at step 1. This loop printed -19.2. So the forward pass is fine, the loss is fine, and the update rule is fine. The important clue is that the value now stored in .grad is not just the derivative of this step’s loss.

That narrows the problem to one place: what was already in .grad before backward() added the new derivative.

And the discrepancy is not random:

-19.2 = -3.2 + (-16.0)

It is this step’s gradient plus the previous step’s. backward() does not assign to .grad; it adds to it. Confirm the mechanism in isolation:

x = torch.tensor(2.0)
target = torch.tensor(6.0)
w = torch.tensor(1.0, requires_grad=True)

for i in range(3):
    loss = (x * w - target) ** 2
    loss.backward()
    print(f"backward {i+1}: w.grad = {w.grad.item()}")

backward 1: w.grad = -16.0
backward 2: w.grad = -32.0
backward 3: w.grad = -48.0

Same w, same loss, three identical gradients of -16, accumulating into -48. This is deliberate on PyTorch’s part, not an oversight. Accumulation lets you intentionally combine gradients from several backward passes before updating, as in gradient accumulation across micro-batches or other workflows that require multiple backward passes.

The cost of that flexibility is that you must decide when the old gradient should be cleared.

The repair is the line we deleted:

w.grad.zero_()

Verify it rather than assuming it. The check is not “does the loss go down again” but “does w.grad at step 1 equal -3.2”. It does. That is the difference between a fix and a coincidence: we predicted a specific number from an independent calculation, and observed it.

It is worth being precise about why lowering the learning rate is the wrong response here. It might make the run look calmer because every corrupted update becomes smaller, but the stale gradients are still present. The update is still being driven by accumulated gradients rather than by the gradient of the current step’s loss.

A symptom can become less visible without its cause being corrected. Before you accept a fix, you should be able to say which quantity it corrected and how you observed the correction.

This is also, concretely, the sort of bug that survives review. A generated training loop that omits zero_grad() runs, prints falling losses for the first step or two, and looks entirely ordinary.

Break it differently: the parameter that stops being a parameter

Here is a second failure in the same family. The update is written as a rebinding rather than an in-place edit:

with torch.no_grad():
    w = w - learning_rate * w.grad     # note: w = w - ..., not w -= ...

Those two forms look interchangeable, and in ordinary Python arithmetic they usually are. Run the loop and print the parameter’s status on each iteration:

x = torch.tensor(2.0)
target = torch.tensor(6.0)
w = torch.tensor(1.0, requires_grad=True)
learning_rate = 0.1

for step in range(3):
    prediction = x * w
    loss = (prediction - target) ** 2

    print(f"step={step} w={w.item():.4f} loss={loss.item():.4f} "
          f"requires_grad={w.requires_grad} is_leaf={w.is_leaf}")

    loss.backward()

    with torch.no_grad():
        w = w - learning_rate * w.grad

step=0 w=1.0000 loss=16.0000 requires_grad=True is_leaf=True
step=1 w=2.6000 loss=0.6400 requires_grad=False is_leaf=True
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

The exception arrives on the second backward(), one iteration after the actual mistake. That delay is typical, and it is why reading the traceback alone tends to send people to the wrong line.

The evidence is in the printed status. At step 0 w requires gradients. At step 1 it does not. Nothing explicitly changed the flag on the original tensor, so something must have replaced the object referred to by w.

That is exactly what happened. Inside torch.no_grad(), the expression w - learning_rate * w.grad creates a new tensor that does not require gradients, and the assignment makes the name w refer to that new tensor. From step 1 onward, the computation no longer depends on a tensor that requires gradients, so backward() has nothing to differentiate with respect to.

The repair is the in-place form, w -= learning_rate * w.grad, which modifies the existing tensor and leaves its identity intact.

There is a tempting alternative that you will see suggested: keep the rebinding and add w.requires_grad_(True) afterwards. The exception disappears and this tiny loop trains. It is still the wrong mental model, because every update replaces w with a fresh trainable leaf tensor.

With one hand-written variable you may never notice. Once an optimizer holds references to tensors, or an nn.Module registers nn.Parameter objects, replacing an object rather than updating the registered one can leave the object being computed with and the object the rest of the system is tracking out of sync. Chapter 4 returns to that failure mode in the context where it becomes genuinely dangerous.

Two chapters from now this distinction will have a name and a mechanism. For the moment, keep the shape of it:

In a training system, it matters which tensor object is being updated, not just which numeric value it currently contains. Once optimizers and modules begin tracking trainable state, replacing an object and modifying the tracked object are no longer equivalent operations.

Loss going down is not the same as learning

Two more experiments, both about the same misconception.

First, vary the learning rate on the original one-parameter problem and watch the first six steps:

def train(lr, steps=6):
    x = torch.tensor(2.0)
    target = torch.tensor(6.0)
    w = torch.tensor(1.0, requires_grad=True)
    for step in range(steps):
        loss = (x * w - target) ** 2
        loss.backward()
        print(f"  step={step} w={w.item():+.4f} loss={loss.item():.4f}")
        with torch.no_grad():
            w -= lr * w.grad
        w.grad.zero_()

for lr in (0.01, 0.1, 0.25, 0.3):
    print(f"lr={lr}")
    train(lr)

Four regimes appear, on identical code and identical data:

lr=0.01                            lr=0.1
  step=0 w=+1.0000 loss=16.0000      step=0 w=+1.0000 loss=16.0000
  step=1 w=+1.1600 loss=13.5424      step=1 w=+2.6000 loss=0.6400
  step=2 w=+1.3072 loss=11.4623      step=2 w=+2.9200 loss=0.0256
  step=3 w=+1.4426 loss=9.7017       step=3 w=+2.9840 loss=0.0010
  step=4 w=+1.5672 loss=8.2115       step=4 w=+2.9968 loss=0.0000
  step=5 w=+1.6818 loss=6.9502       step=5 w=+2.9994 loss=0.0000

lr=0.25                            lr=0.3
  step=0 w=+1.0000 loss=16.0000      step=0 w=+1.0000 loss=16.0000
  step=1 w=+5.0000 loss=16.0000      step=1 w=+5.8000 loss=31.3600
  step=2 w=+1.0000 loss=16.0000      step=2 w=-0.9200 loss=61.4656
  step=3 w=+5.0000 loss=16.0000      step=3 w=+8.4880 loss=120.4726
  step=4 w=+1.0000 loss=16.0000      step=4 w=-4.6832 loss=236.1264
  step=5 w=+5.0000 loss=16.0000      step=5 w=+13.7565 loss=462.8078

At 0.01 the direction is right and the model is learning, just slowly; given enough steps it arrives. At 0.1 it converges. At 0.25 it steps clean over the minimum and lands the same distance away on the other side, then steps back, forever: a loss that is perfectly stable and perfectly useless. At 0.3 it overshoots by more than it started with and diverges. Nothing here is a bug. The same correct code produces all four behaviors, and only the printed numbers distinguish them.

That is the first version of the point. Here is a sharper one.

Give the model a second parameter, a bias, and train it on a single example:

x = torch.tensor(2.0)
target = torch.tensor(6.0)

w = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)

learning_rate = 0.05

for step in range(60):
    prediction = w * x + b
    loss = (prediction - target) ** 2

    loss.backward()

    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad

    w.grad.zero_()
    b.grad.zero_()

print("final training loss:", ((w * x + b - target) ** 2).item())
print("w =", round(w.item(), 4), " b =", round(b.item(), 4))

for value in (1.0, 3.0, 5.0):
    xt = torch.tensor(value)
    print(f"x={value}: prediction={(w * xt + b).item():.4f}  truth={3 * value}")

final training loss: 0.0
w = 2.6  b = 0.8

x=1.0: prediction=3.4000  truth=3.0
x=3.0: prediction=8.6000  truth=9.0
x=5.0: prediction=13.8000  truth=15.0

The training loss is zero. Not small: zero. And the model is wrong everywhere except the one point it was shown. It did not find y = 3x; it found y = 2.6x + 0.8, which happens to pass exactly through (2, 6). So does y = x + 4, and so does every other line through that point. With two parameters and one example, the problem is underdetermined, and gradient descent has no reason to prefer the answer we had in mind. It solved the problem we actually posed.

This deserves to be stated plainly, because it is the assumption that quietly underlies most wasted debugging time:

A falling training loss is evidence that the optimization process is reducing the objective you gave it. It is not evidence that the model has learned the behavior you actually care about. Those are separate claims, and they need separate evidence.

The repair is not in the loop. It is in the problem. Give the model more than one example, and average the loss across them:

x = torch.tensor([1.0, 2.0, 3.0, 4.0])
target = torch.tensor([3.0, 6.0, 9.0, 12.0])

w = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)

learning_rate = 0.05

for step in range(500):
    prediction = w * x + b
    loss = ((prediction - target) ** 2).mean()

    loss.backward()

    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad

    w.grad.zero_()
    b.grad.zero_()

    if step in (0, 9, 99, 499):
        print(f"step={step:3d} loss={loss.item():.6f} "
              f"w={w.item():.4f} b={b.item():.4f}")

held_out = torch.tensor(10.0)
print(f"held-out x=10 -> prediction {(w * held_out + b).item():.4f}, truth 30.0")

step=  0 loss=30.000000 w=2.5000 b=0.5000
step=  9 loss=0.047325 w=2.8217 b=0.5243
step= 99 loss=0.003134 w=2.9541 b=0.1349
step=499 loss=0.000000 w=2.9999 b=0.0003

held-out x=10 -> prediction 29.9992, truth 30.0

Four points instead of one, and w converges to 3 and b to 0. More importantly, the model also produces the expected result at x = 10, an extrapolation point it never saw during training.

That is stronger evidence than training loss alone, although one successful unseen prediction is not enough to establish generalization. In a real experiment we would evaluate across held-out data and compare against appropriate baselines. Chapter 14 turns that idea into a discipline.

Two mechanical details in that loop point forward. prediction is now a tensor of four values rather than one, because multiplying a scalar parameter by a four-element input produces four results; Chapter 2 is about exactly what rules govern that.

And .mean() collapses the four per-example losses into the single scalar objective we want to optimize. That scalar matters because a bare loss.backward() can infer the starting gradient automatically for a scalar output. PyTorch can also backpropagate from a non-scalar tensor, but then you must explicitly supply the gradient vector that defines the Jacobian-vector product. Chapter 3 takes that distinction apart.

What PyTorch actually contributed

Look back at what we have written. We chose the model, chose the loss, chose the update rule, and decided when to stop. We derived the gradient ourselves once so that we had an independent result to check; PyTorch then automated that differentiation for us.

What it supplied was bookkeeping, and the bookkeeping is the part that does not scale by hand:

  • a numeric container, the tensor, that carries not only values but a shape, a dtype, a device, and a memory layout;
  • a record of which operations produced which values, built automatically as the forward pass executes;
  • differentiation of that record, so the chain rule is applied for you no matter how long the chain gets;
  • later, the organization of parameter state, the movement of work onto accelerators, and the machinery that feeds data in.

With one parameter, differentiating by hand was a two-line exercise. With a transformer it is not an exercise at all. That is the trade the framework offers, and it is a good one, as long as you remember what was traded. The framework automates the calculation. It does not automate the understanding, and when the calculation produces something unexpected, only the understanding helps.

Which is why, when you eventually meet this:

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

it should not read as a neural-network incantation. nn.Linear(784, 128) contains a weight and a bias that play the same autograd role that w played in this chapter: they are trainable leaf tensors whose gradients can accumulate during backward().

There is one important additional mechanism. Inside an nn.Module, those values are nn.Parameter objects, which means PyTorch registers them as model parameters so that operations such as parameters() and state_dict() can discover them. Chapter 5 takes that machinery apart.

Using AI on this loop

You will ask an assistant to write training loops, and it will usually produce correct ones. The failures in this chapter are precisely the kind that can survive a quick inspection: a missing gradient clear, an update that replaces the tensor being trained, or a loop that successfully minimizes the wrong objective.

Some fail silently. Others raise an exception only after the line that actually caused the problem has already executed. In both cases, the visible symptom may point somewhere other than the cause.

The way to get real value from an assistant here is to ask it to commit to predictions you can check, before you let it change anything. Try this on the broken loop from earlier, the one with w.grad.zero_() removed:

Here is a PyTorch training loop. Do not modify it and do not tell me how to fix it yet.

Trace one full iteration. List each named tensor in the training loop and the important intermediate tensors that participate in the path from w to loss. For each one, state whether it requires gradients, whether it is a leaf, and whether its .grad field should be populated after backward() returns.

Then predict the exact numeric value that w.grad will hold at step 0 and at step 1, and show the arithmetic you used.

Finally, list the print statements I could add to check each of your predictions.

Then run it and compare. If the predictions match what you observe, you have evidence that your model of this part of the program is correct. If they do not, the disagreement is valuable because it tells you which assumption needs to be investigated next. Either way you have gained something a request to “fix this” would not have given you.

Ask for traces, predictions, comparisons and experiments. Ask for a fix once you know what is broken. A proposed fix you cannot check is a guess with better grammar than yours.

What you should now be able to answer

Here is a complete training program, small enough to hold in your head:

import torch

x = torch.tensor([1.0, 2.0, 3.0, 4.0])
target = torch.tensor([3.0, 6.0, 9.0, 12.0])

w = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)

for step in range(500):
    prediction = w * x + b
    loss = ((prediction - target) ** 2).mean()
    loss.backward()
    with torch.no_grad():
        w -= 0.05 * w.grad
        b -= 0.05 * b.grad
    w.grad.zero_()
    b.grad.zero_()

You should be able to answer each of the following without running it.

What are the learnable values? w and b. They are the tensors created with requires_grad=True, and they are the only things the update statements modify.

What produced the prediction? w * x + b, a multiplication and an addition, both recorded because w and b require gradients.

What produced the loss? A subtraction, a squaring, and a mean, reducing four per-example errors to one scalar.

Which operations should be differentiable? Every operation on the path from w and b to loss. Nothing inside the torch.no_grad() block.

What does backward() actually do? It traverses the recorded operations from loss backwards, applies the chain rule, and accumulates the result into w.grad and b.grad. It changes no parameter values.

Which .grad fields should be populated? w.grad and b.grad. x.grad and target.grad stay None because those tensors do not require gradients. Intermediate gradients are computed as needed during backpropagation, but PyTorch does not retain them in .grad by default.

What operation actually changes the parameters? The two in-place subtractions inside the no_grad block. Remove them and backward() will run five hundred times and train nothing.

What evidence would show that learning is occurring? In this toy problem, w should move toward 3, b toward 0, the training objective should fall, and predictions on unseen inputs should approach the rule we intended to learn. The shrinking gradient is useful evidence here because we know the shape of this particular loss surface; later, small gradients can mean several very different things.

If a parameter does not change, what should I inspect first? In order: is loss.requires_grad still True? Is w.grad None, or zero, immediately after backward()? Does w still report requires_grad=True and is_leaf=True, or has something replaced it? And is the update statement actually inside the loop?

That last question is the one this book is really about. The answer is a sequence of things to look at, each of which produces evidence, and each of which rules something out. As the systems get larger the list gets longer and the evidence gets harder to obtain, but the method does not change.

Make the hidden structure visible before guessing.

Exercises

These are worth running rather than reading, and they map onto the notebook that accompanies this chapter.

  1. Reproduce the gradient by hand. For w = 1.5, x = 2, target = 6, compute dloss/dw on paper, then check it against w.grad. Repeat for the two-parameter model with a bias, deriving dloss/db as well.

  2. Watch accumulation directly. Call backward() three times on the same loss without clearing, printing w.grad each time. Then replace w.grad.zero_() with w.grad = None and confirm the loop still trains. Explain why both work.

  3. Confirm that the parameter moved. Before each update, save before = w.detach().clone(). After the update, print (w.detach() - before).item() and compare it with -learning_rate * gradient. This one-line check is the fastest way to answer “did the optimizer do anything” in real code.

  4. Break gradient tracking on purpose, twice. First create w without requires_grad=True. Then restore it and instead insert prediction = prediction.detach() before the loss. Predict the error message in each case before running. You will find that both produce the identical RuntimeError, despite having different causes, which is a small and useful lesson: the symptom does not always identify the cause, and the thing that distinguishes them is evidence you go and collect.

  5. Find the boundary. For the one-parameter problem, the learning rate that oscillates without converging is 0.25. Derive that number from the update rule, then confirm it empirically. Establish what happens just below and just above it.

  6. Make the model identifiable. Start from the w-and-b model trained on a single point and add a second example with a different x value. Explain why two distinct points determine a unique straight line while one point does not. Then try duplicating the original example instead of adding a new x value. Why does having more rows of data not necessarily give you more information?

Next: the tensor

We have been treating torch.tensor(2.0) as a container for a number, which was enough while every value was a scalar. It stops being enough almost immediately. In the last experiment x held four values and prediction did too, and the multiplication that produced it was quietly following rules we never stated.

Those rules account for a remarkable amount of practical PyTorch debugging time. A tensor carries a shape, a dtype, a device and a memory layout, and each of those can produce a failure with its own characteristic symptom: an operation that refuses to run, an operation that runs and silently produces something the wrong size, a view() that fails after a permute(), a broadcast that turns a [32, 1] and a [1, 32] into a [32, 32] nobody wanted.

The next chapter makes shape reasoning operational: how to read a shape, how to derive the shape an operation should produce, how to state your assumptions in code, and how to find the line where a dimension first went wrong.

The tensor.