Table of Contents
What autograd is for
PyTorch autograd is the system that automatically computes gradients for you. In deep learning you usually define a model that produces a scalar loss value, then you need the gradients of that loss with respect to the model parameters so an optimizer can update them. Autograd does this by tracking how tensors are computed and then applying the chain rule to propagate derivatives backward.
Autograd only computes gradients for tensors that have requires_grad=True, and gradients are only created by calling backward() on a scalar loss or on a tensor paired with an explicit gradient.
The computational graph in PyTorch
When you perform operations on tensors, PyTorch builds a dynamic computational graph. Dynamic means the graph is created as your Python code runs, operation by operation. Each result tensor can carry a reference to a function that created it, accessible via attributes like grad_fn. Leaf tensors are tensors that you created directly, not the result of an operation. Model parameters are typically leaf tensors with requires_grad=True.
A typical pattern is that your inputs might not require gradients, but your parameters do. When you compute predictions and then a loss, the loss tensor sits at the end of the graph. Calling loss.backward() walks the graph in reverse and accumulates gradients into the .grad fields of leaf tensors that require gradients.
A minimal example of gradient tracking
Consider a simple scalar computation:
import torch
x = torch.tensor(2.0, requires_grad=True)
y = 3 * x**2 + 1
y.backward()
print(x.grad) # dy/dx
Here autograd records multiplication, power, and addition. When you call backward(), it computes $\frac{dy}{dx}$ and stores it in x.grad. For this example, $y = 3x^2 + 1$, so $\frac{dy}{dx} = 6x$, which at $x=2$ is $12$.
Gradients produced by backward() are accumulated into .grad by addition. If you call backward() again without clearing .grad, you will add new gradients on top of the old ones.
Vector outputs and providing gradients
backward() is most straightforward when the tensor you call it on is a scalar, such as a loss. If you have a non scalar tensor, PyTorch needs to know what combination of its elements you want to differentiate. Mathematically, backpropagation computes vector Jacobian products. In practice, you provide an upstream gradient of the same shape.
x = torch.randn(3, requires_grad=True)
y = x * x # shape (3,)
y.backward(torch.ones_like(y)) # sums gradients as if loss = y.sum()
print(x.grad) # should be 2*x
Passing torch.ones_like(y) makes this equivalent to differentiating $y_1 + y_2 + y_3$.
If you call backward() on a non scalar tensor without specifying gradient=..., PyTorch will raise an error. Losses are typically reduced to scalars using .mean() or .sum() to avoid this.
Leaf tensors, `.grad`, and `grad_fn`
Understanding where gradients appear helps with debugging. The .grad field is populated for leaf tensors that require gradients. Intermediate tensors usually have grad_fn but do not keep .grad by default.
w = torch.randn(5, requires_grad=True) # leaf
x = torch.randn(5) # does not require grad
y = (w * x).sum() # scalar
print(y.grad_fn) # shows a backward function node
y.backward()
print(w.grad) # gradients exist here
print(y.grad) # usually None, not a leaf
If you ever need gradients of a non leaf tensor for inspection, you can ask PyTorch to retain them with tensor.retain_grad() before calling backward(). Use this sparingly since it increases memory use.
Detaching tensors and stopping gradients
Sometimes you want to use a tensor value but prevent gradients from flowing back through it. You can do this by detaching the tensor from the graph.
a = torch.randn(3, requires_grad=True)
b = (a * 2).detach() # shares storage but has no grad history
c = (b * 3).sum()
c.backward() # no gradient flows to a
print(a.grad) # NoneDetaching is common when logging values, creating targets from model outputs, or implementing algorithms that require a stop gradient operation.
detach() breaks the gradient path. If you detach something by accident, your model may stop learning because parameters receive no gradients.
Turning autograd off for inference
During evaluation or inference you typically do not want to build graphs at all. Disabling autograd reduces memory usage and usually speeds up computation. Use torch.no_grad() as a context manager.
model.eval()
with torch.no_grad():
preds = model(x)
There is also torch.set_grad_enabled(flag) which is useful when you want code that can run in both training and evaluation modes depending on a boolean.
Forgetting torch.no_grad() during inference can cause unnecessary memory growth because PyTorch keeps graph history for operations on tensors that require gradients.
Backward graphs, `retain_graph`, and multiple backward passes
By default, PyTorch frees the graph buffers after backward() to save memory. If you need to call backward() more than once on the same graph, you must keep it:
loss.backward(retain_graph=True)
loss2.backward()This is only necessary in specific cases, such as certain multi objective setups or when computing multiple gradient based quantities from the same forward pass.
Using retain_graph=True increases memory usage. Prefer recomputing the forward pass unless you have a clear reason to retain the graph.
Higher order gradients with `create_graph`
Some techniques need gradients of gradients, such as certain meta learning methods or gradient penalty terms. You can tell autograd to construct a graph during the backward pass with create_graph=True.
x = torch.tensor(2.0, requires_grad=True)
y = x**3
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]This is an advanced feature, but it is useful to recognize what it does: it allows differentiation through the gradient computation itself.
`backward()` vs `torch.autograd.grad`
loss.backward() is convenient and fills .grad on leaf tensors. torch.autograd.grad returns gradients directly without modifying .grad unless you assign them yourself. It is useful when you want gradients for specific tensors, when you want to avoid accumulation in .grad, or when computing higher order derivatives.
w = torch.randn(3, requires_grad=True)
loss = (w**2).sum()
g = torch.autograd.grad(loss, w)[0]
Here g is the gradient tensor, and w.grad remains None unless you call backward().
In place operations and autograd safety
Autograd needs saved intermediate values to compute derivatives. Some in place operations can overwrite values that autograd expects to reuse, causing runtime errors or incorrect gradients. PyTorch often detects this and raises an error mentioning that a variable needed for gradient computation has been modified in place.
Avoid in place operations on tensors that require gradients, especially when they are part of the forward pass that feeds into the loss. If you see an in place modification error, replace operations like x += ... with x = x + ... in the relevant part of the computation.
The chain rule viewpoint you should keep in mind
Backpropagation is an application of the chain rule. If $L$ is the loss and parameters are $\theta$, training needs $\nabla_{\theta} L$. Autograd does this by composing local derivatives along the graph.
The essential workflow is: build the loss as a scalar, call loss.backward(), read parameter gradients from param.grad, then update parameters and clear gradients before the next step.