Understanding the engine explains in-place errors, memory use in backward, and how to compute Hessians or meta-gradients.
During the forward pass, every differentiable op creates a node (grad_fn) holding references to the tensors it needs for its derivative (the saved tensors) and edges to the nodes of its inputs (grad_fn.next_functions). User-created tensors with requires_grad=True are leaves; they get an AccumulateGrad node that writes into .grad. Non-leaf tensors do not keep .grad unless you call retain_grad().
backward() runs a topological traversal from the output, computing vector-Jacobian products node by node, so the cost of a full gradient is a small constant times the forward cost. Saved tensors are freed as soon as their node has run, which is why the default graph is single-use.
In-place safety comes from version counters. Every tensor has a _version that in-place ops increment. When a node is executed, it checks that each saved tensor still has the version it saw; if not, you get one of the variables needed for gradient computation has been modified by an inplace operation.
torch.autograd.grad(outputs, inputs, create_graph=True) returns gradients that are themselves part of a graph, so you can differentiate again: Hessian-vector products, gradient penalties and meta-learning. For a functional, composable API, torch.func provides grad, vmap, jacrev and hessian.
Saved tensors dominate training memory; torch.autograd.graph.saved_tensors_hooks lets you compress or offload them.
Example
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x * x
z = y.sin()
print(z.grad_fn) # <SinBackward0 object at ...>
print(z.grad_fn.next_functions) # ((<MulBackward0 object at ...>, 0),)
print(x.is_leaf, y.is_leaf) # True False
(g,) = torch.autograd.grad(z, x, create_graph=True) # dz/dx = 2x cos(x^2)
(h,) = torch.autograd.grad(g, x) # second derivative
print(round(g.item(), 4), round(h.item(), 4)) # -2.6146 10.8016
a = torch.ones(3, requires_grad=True)
b = a.exp() # exp saves its output for backward
b.add_(1) # in-place bump of b's version counter
try:
b.sum().backward()
except RuntimeError as e:
print(str(e).split(":")[0])
# one of the variables needed for gradient computation has been modified by an inplace operationRemember
- The graph is grad_fn nodes with saved tensors and edges to inputs
- Version counters detect in-place edits of saved tensors
create_graph=Trueenables higher-order derivatives
Watch out for
- Using
retain_graph=Trueto silence errors, which hides a design problem and leaks memory.