Understanding PyTorch Autograd by Writing a Custom Training Loop

You have almost certainly typed loss.backward() without knowing what it does. Most people do, for a long time. It sits there in the middle of the training loop looking like an instruction to a black box, and as long as the model trains, there is no reason to ask questions. Then one day you get an error saying you tried to go backward through the graph a second time, and suddenly the black box matters a great deal. This tutorial opens it up. You will watch PyTorch build a record of your calculations as they happen, walk that record backwards, and then throw it away. By the end, the rules you have been following on faith will simply look like consequences of a design you understand.

What Automatic Differentiation Solves

Training a model means answering one question, repeatedly, for every adjustable number it contains: if I nudge this weight slightly, does the error go up or down, and by how much? The answer is a derivative, and there are three ways to get one.

You could work it out symbolically with algebra, which is exact but becomes unmanageable once the expression is a hundred layers deep. You could estimate it numerically by trying a slightly different value and seeing what changes, which is simple but slow and imprecise. Or you could record each operation as it happens, then replay the recording backwards applying a known reversal rule at each step.

That third approach is automatic differentiation, and it is what autograd does. It is exact like the symbolic method and fast like the numerical one, because every primitive operation in PyTorch ships with its derivative already written by hand.

Why the Graph Is Built While Your Code Runs

The most consequential design decision in PyTorch is that this recording is made during execution rather than declared in advance. Frameworks that require you to define the full structure up front can optimise it more aggressively, but they take your Python away from you in the process.

This is the main reason PyTorch took over research, and why most new models are released as PyTorch code first. You can put an if statement inside a model and it simply works.

The Three Rules That Explain Every Gradient Error

Nearly every confusing autograd error you will meet follows from three facts about how the recording is kept. Everything demonstrated below is a consequence of one of them.

The recording is discarded once it has been replayed. Backward frees the intermediate values it needed, because holding them costs memory and almost nobody wants them twice. This is why calling backward a second time fails.

Results accumulate rather than replace. Each backward pass adds its findings to whatever is already stored on a parameter. This is why something has to clear them between steps.

Only tensors you created directly receive results. Model parameters get gradients. Intermediate values calculated along the way do not, because no optimizer would ever update them.

Hold those three and the error messages stop being mysterious. The rest of this article shows each one happening in printed output.

Step 1: Watch a Tensor Start Tracking

Set requires_grad=True and PyTorch begins recording.

import torch

x = torch.tensor([2.0], requires_grad=True)
y = x ** 3

print("x        :", x)
print("y = x**3 :", y)
print("y.grad_fn:", y.grad_fn)

y.backward()
print("dy/dx at x=2:", x.grad.item())

x        : tensor([2.], requires_grad=True)
y = x**3 : tensor([8.], grad_fn=<PowBackward0>)
y.grad_fn: <PowBackward0 object at 0x7fa40d51b280>
dy/dx at x=2: 12.0

Three things happened worth naming.

y carries a grad_fn. The tensor holds the value 8.0 and also a reference to PowBackward0, the object that knows how to reverse a power operation. Every tensor produced from a tracked tensor gets one of these.

backward() walked the graph in reverse. It started at y, found PowBackward0, applied the derivative rule for powers, and deposited the result on x.

The answer is exact. The derivative of x cubed is 3 times x squared, which at x equals 2 gives 12. Autograd returned exactly 12.0, not an approximation from sampling nearby points.

Step 2: See the Graph Structure

The graph is a chain of backward functions, and you can inspect it directly.

p = torch.tensor([4.0], requires_grad=True)
q = p * 2
r = q + 1

print("r.grad_fn               :", r.grad_fn)
print("r.grad_fn.next_functions:", r.grad_fn.next_functions)

r.grad_fn               : <AddBackward0 object at 0x7fa3eeaeebc0>
r.grad_fn.next_functions: ((<MulBackward0 object at 0x7fa3eeaeec80>, 0), (None, 0))

Read that backwards, which is how autograd reads it. The final operation was an addition, so r holds AddBackward0. That node points to MulBackward0, the multiplication that produced q. The None is the constant 1, which needs no gradient because it is not a tracked tensor.

In a real network this chain runs through every layer, every activation, and every arithmetic operation between the input and the loss. Calling .backward() on the loss traverses the whole thing.

Step 3: Prove That Gradients Accumulate

This is the behaviour behind the most repeated rule in PyTorch, and it is easier to accept once you have watched it.

a = torch.tensor([1.0], requires_grad=True)
for step in range(1, 4):
    out = (a * 3).sum()
    out.backward()
    print(f"after backward #{step}: a.grad = {a.grad.item()}")

after backward #1: a.grad = 3.0
after backward #2: a.grad = 6.0
after backward #3: a.grad = 9.0

The correct gradient is 3.0 every time, because the derivative of 3a with respect to a is 3 regardless of how many times you ask. What changes is the stored value, because each call adds to what is already there.

Now the same loop with a reset:

a2 = torch.tensor([1.0], requires_grad=True)
for step in range(1, 4):
    if a2.grad is not None:
        a2.grad.zero_()
    out = (a2 * 3).sum()
    out.backward()
    print(f"after backward #{step}: a2.grad = {a2.grad.item()}")

after backward #1: a2.grad = 3.0
after backward #2: a2.grad = 3.0
after backward #3: a2.grad = 3.0

That is what optimizer.zero_grad() does across every parameter in your model at once.

This design is deliberate rather than an oversight. Accumulation is what makes gradient accumulation across several small batches possible when a single large batch will not fit in memory, a routine technique when training large models on limited hardware. PyTorch's position is that you should opt out of accumulation explicitly rather than have it opted out for you.

Step 4: Leaf Tensors and Why .grad Is Often None

Not every tensor receives a gradient, and the distinction confuses almost everyone at first.

leaf = torch.tensor([3.0], requires_grad=True)
mid = leaf * 2
final = mid.sum()
final.backward()

print("leaf.is_leaf:", leaf.is_leaf, "| leaf.grad:", leaf.grad.item())
print("mid.is_leaf :", mid.is_leaf,  "| mid.grad :", mid.grad)

leaf.is_leaf: True | leaf.grad: 2.0
mid.is_leaf : False | mid.grad : None

PyTorch also prints a warning on that second line:

UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being
accessed. Its .grad attribute won't be populated during autograd.backward(). If you
indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad()
on the non-leaf Tensor.

A leaf tensor is one you created directly rather than one that came out of an operation. Model parameters are leaves. Activations and intermediate results are not.

Autograd only stores gradients on leaves because those are the only tensors an optimizer updates. Keeping gradients for every intermediate value would consume memory for data nobody uses. If you genuinely need an intermediate gradient, usually for debugging or for techniques like saliency mapping, call .retain_grad() on that tensor before the backward pass.

Step 5: Two Ways to Stop Tracking

torch.no_grad() and .detach() both disable gradient tracking, and they are not interchangeable.

m = torch.tensor([5.0], requires_grad=True)

n = m * 2
print("n.requires_grad          :", n.requires_grad)
print("n.detach().requires_grad :", n.detach().requires_grad)

with torch.no_grad():
    o = m * 2
print("inside no_grad()         :", o.requires_grad)

n.requires_grad          : True
n.detach().requires_grad : False
inside no_grad()         : False

The .item() row matters more than it looks. Writing running_loss += loss instead of running_loss += loss.item() keeps every batch's entire computation graph alive in memory, and by the end of an epoch the process has consumed several gigabytes for a number you only wanted to print.

Step 6: A Complete Training Loop Annotated by Graph State

Here is the standard loop with a note on what autograd is doing at each line.

from torch import nn

model = nn.Linear(4, 2)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
x, y = torch.randn(8, 4), torch.randint(0, 2, (8,))

for epoch in range(3):
    pred = model(x)              # graph is built here, node by node
    loss = loss_fn(pred, y)      # graph now runs from x through to loss

    optimizer.zero_grad()        # every parameter .grad is set to zero
    loss.backward()              # graph traversed in reverse, .grad populated, graph freed
    optimizer.step()             # parameters updated using .grad, no graph involved

    print(f"epoch {epoch} | loss {loss.item():.4f}")

The ordering follows from the mechanics. Gradients must be cleared before backward() because backward() adds to them. step() must come after backward() because it reads the values backward() wrote. And the forward pass must be inside the loop because the graph is destroyed on every backward call.

Break It On Purpose

Two experiments, each taking under a minute, that make the rules stick.

x = torch.tensor([2.0], requires_grad=True)
y = x * 3
y.backward()
y.backward()   # RuntimeError

RuntimeError: Trying to backward through the graph a second time
(or directly access saved tensors after they have already been freed).

w = torch.zeros(1, requires_grad=True)
w -= 0.1   # RuntimeError

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

Error Decoder

Every message below was produced by running the broken code on PyTorch 2.13.0.

Why the Dynamic Graph Design Matters

PyTorch rebuilds the graph on every forward pass. That sounds wasteful, and it has a real cost, but it buys something valuable: your model is ordinary Python.

You can put an if statement in forward and take different paths for different inputs. You can loop a variable number of times based on sequence length. You can insert a print in the middle of a forward pass and see actual tensor values, or set a breakpoint and step through with a standard Python debugger. None of that requires special framework support, because there is no separate compilation step deciding the structure in advance.

This is the main reason PyTorch became dominant in research, and why most new models are released as PyTorch code first. When you need the performance of a static graph, torch.compile() traces your model and optimises it, so you develop with the flexible version and switch on the fast one when the code is settled.

Conclusion

Autograd stops feeling like magic once you accept three facts. The graph is built as your code runs and destroyed as soon as you traverse it backwards. Gradients accumulate on leaf tensors until something clears them. Tracking can be switched off for a whole block with torch.no_grad() or for one tensor with .detach().

Every gradient rule you have been following falls out of those facts. Clear before backward because backward adds. Rebuild the forward pass each iteration because the graph is gone. Wrap evaluation in no_grad() because you are not going to differentiate it and building the graph wastes memory. Call .item() when logging because otherwise you keep the whole graph alive for a single float.

Spend ten minutes in a REPL printing grad_fn, is_leaf, and .grad on tensors as you build small expressions. That direct experimentation does more than any amount of reading, because the objects are inspectable and the behaviour is completely deterministic.

FAQS