Autograd

PyTorch autograd automatically computes gradients of scalar outputs with respect to tensor inputs. It powers all neural network training through reverse-mode automatic differentiation.

Recipe

Quick-reference recipe card - copy-paste ready.

import torch
 
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print(x.grad)  # dy/dx at x=2: 2*2+3 = 7
 
# Inference: no gradient tracking
with torch.no_grad():
    pred = model(x)

When to reach for this:

• Understanding why loss.backward() and optimizer.step() work.
• Debugging gradient flow in custom layers or loss functions.
• Implementing custom autograd functions for non-standard operations.
• Disabling gradients during inference or evaluation.

Working Example

"""autograd.py - gradient computation and gradient checking."""
import torch
import torch.nn as nn
 
# Simple computation graph
w = torch.tensor([1.0, 2.0], requires_grad=True)
b = torch.tensor(0.5, requires_grad=True)
x = torch.tensor([3.0, 4.0])
 
# Forward: y = w . x + b, loss = y^2
y = (w * x).sum() + b
loss = y ** 2
loss.backward()
 
print("w.grad:", w.grad)  # d(loss)/dw
print("b.grad:", b.grad)  # d(loss)/db
 
# Model with autograd
model = nn.Linear(3, 1)
x_batch = torch.randn(16, 3)
target = torch.randn(16, 1)
 
model.zero_grad(set_to_none=True)
pred = model(x_batch)
loss = nn.functional.mse_loss(pred, target)
loss.backward()
 
for name, param in model.named_parameters():
    print(f"{name}: grad norm = {param.grad.norm().item():.4f}")

What this demonstrates:

• Leaf tensors (w, b) accumulate gradients in .grad.
• loss.backward() traverses the graph in reverse.
• Model parameters receive gradients automatically.
• zero_grad(set_to_none=True) clears gradients efficiently.

Deep Dive

How It Works

• Each tensor operation records inputs and the backward function on a computation graph.
• backward() applies the chain rule from the output back to leaves.
• Only tensors with requires_grad=True accumulate .grad.
• Non-leaf tensor gradients are freed after backward unless retain_graph=True.
• torch.inference_mode() is stricter and faster than no_grad() for inference.

Gradient Controls

Python Notes

# Custom autograd function
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)
 
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

Gotchas

• Forgetting zero_grad() - gradients accumulate across steps. Fix: call optimizer.zero_grad(set_to_none=True) before each backward.
• Calling backward twice without retain_graph - graph is freed after first backward. Fix: loss.backward(retain_graph=True) or separate forward passes.
• In-place operations on tensors needing grad - breaks autograd tracking. Fix: avoid x.add_(1) on tracked tensors; use x = x + 1.
• GPU memory from saved graphs - keeping references to loss tensors retains the graph. Fix: use .item() for logging; del loss after backward.
• Gradients on integer tensors - only floating/complex tensors support autograd. Fix: cast to float before operations needing gradients.
• Mixing numpy and torch - in-place numpy changes on shared memory corrupt gradients. Fix: .detach().cpu().numpy() for export.

Alternatives

FAQs

for name, p in model.named_parameters():
    if p.grad is None:
        print(f"no grad: {name}")

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

loss.backward(create_graph=True)
grad_grad = torch.autograd.grad(loss, model.parameters(), create_graph=True)

with torch.no_grad():
    features = frozen_encoder(x)
output = trainable_head(features)

Related

• PyTorch Basics - tensor fundamentals
• Training Loops - backward in practice
• Building Models with nn.Module - nn.Parameter
• Transfer Learning & Fine-Tuning - freezing layers
• GPUs & Mixed Precision - autocast and grad scaling

Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+.