PyTorch Autograd Explained

What is autograd?#

Autograd is PyTorch's automatic differentiation engine. As you compute with tensors, PyTorch records every operation in a computation graph. Calling .backward() walks that graph backwards and computes the gradient of the result with respect to every tensor that needs one — the backpropagation you'd otherwise write by hand.

A first example#

Let's compute the gradient of loss = (3w − 12)² at w = 2.

By hand: d(loss)/dw = 2 · (3w − 12) · 3 = 2 · (6 − 12) · 3 = −36.

import torch

w = torch.tensor(2.0, requires_grad=True)   # track gradients for w
loss = (3 * w - 12) ** 2                     # PyTorch records each operation
loss.backward()                              # compute d(loss)/dw
print(w.grad)

tensor(-36.)

Autograd matched the calculus exactly.

How it works#

1. requires_grad=True marks a tensor as something to learn (model parameters have this set automatically).
2. During the forward pass, each operation creates a node in the graph and remembers how to compute its local derivative.
3. loss.backward() applies the chain rule from loss back to every leaf tensor with requires_grad=True.
4. The results are stored in each tensor's .grad attribute.

Inspect the graph:

x = torch.tensor(3.0, requires_grad=True)
y = x * x + 2 * x
print(y.grad_fn)      # the operation that produced y
y.backward()
print(x.grad)         # dy/dx = 2x + 2 = 8 at x = 3

<AddBackward0 object at 0x...>
tensor(8.)

Using gradients to learn: manual gradient descent#

w = torch.tensor(2.0, requires_grad=True)

for step in range(50):
    loss = (3 * w - 12) ** 2
    loss.backward()
    with torch.no_grad():          # update without recording it in the graph
        w -= 0.01 * w.grad
    w.grad.zero_()                 # reset the gradient for the next step

print(round(w.item(), 3))          # approaches 4.0, where 3w - 12 = 0

4.0

Two essential details:

• torch.no_grad() — the weight update itself must not become part of the graph.
• w.grad.zero_() — gradients accumulate by default (they're added, not replaced). Forgetting to reset them is a classic bug. With an optimizer, this becomes optimizer.zero_grad().

Turning gradient tracking off#

Tracking gradients costs memory and time. Turn it off when you don't need it:

Gradients of non-scalar outputs#

.backward() works directly on a scalar (like a mean loss). For a tensor output you either reduce it first (out.sum().backward()) or pass a gradient argument — in training, you almost always backpropagate a scalar loss.

Freezing part of a model#

import torch.nn as nn

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2))

for param in model[0].parameters():     # freeze the first layer
    param.requires_grad_(False)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(trainable)                         # only the last layer: 32*2 + 2 = 66

66

Freezing is the foundation of transfer learning and LoRA-style fine-tuning: train a small part, keep the rest fixed.

Common mistakes#

• Not zeroing gradients between steps, so they pile up.
• Updating weights outside torch.no_grad(), which corrupts the graph or raises an error.
• Calling .backward() twice on the same graph — the graph is freed after the first call unless you pass retain_graph=True (rarely needed).
• Evaluating with gradients on, wasting memory — use inference_mode().

Key takeaways#

• Autograd records operations and computes gradients with .backward().
• Gradients land in .grad and accumulate — reset them each step.
• Use no_grad/inference_mode for updates and evaluation, and requires_grad_(False) to freeze layers.