Every PyTorch tensor that participates in a computation silently carries a history of operations leading to it. When you call .backward(), PyTorch walks that history in reverse, accumulating gradients at every leaf tensor. But the graph is built only once and destroyed after backpropagation—unless you call retain_graph=True. Understanding what the graph is, how to visualize it, and when to break it with detach() separates confident debugging from hours of frustration.
The Computation Graph: A Hidden Data Structure
During a forward pass, every operation (add, multiply, ReLU, etc.) records itself as a node in a graph. The edges connect inputs to outputs. At the end of forward, you have a DAG (directed acyclic graph) from inputs to loss.
import torch
# Simple computation
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 # power operation
z = y * 3 # multiply operation
loss = z.sum() # sum operation
# Inspect the graph
print(f"loss.grad_fn: {loss.grad_fn}")
print(f"loss.grad_fn.next_functions: {loss.grad_fn.next_functions}")
# Traverse the graph
print("\n--- Graph Structure ---")
print(f"loss -> {loss.grad_fn}") # SumBackward0
print(f" -> {loss.grad_fn.next_functions[0][0]}") # MulBackward0
print(f" -> {loss.grad_fn.next_functions[0][0].next_functions}") # PowerBackward0, etc.
# Leaf tensor: no grad_fn
print(f"\nx.grad_fn: {x.grad_fn}") # None — x is a leaf
Output:
loss.grad_fn: SumBackward0()
loss.grad_fn.next_functions: ((MulBackward0(next_functions=(...)), 0),)
--- Graph Structure ---
loss -> SumBackward0()
-> (MulBackward0(next_functions=(...)), 0)
-> ((AccumulateGrad(), 0), (PowBackward0(next_functions=(...)), 1))
x.grad_fn: None
The graph is a tree of operations. Each node knows how to compute gradients via its backward() method. SumBackward0 knows how to backprop through a sum, MulBackward0 through a multiply, and so on. Leaf tensors (created with torch.tensor() or parameters) have no grad_fn—they're the end points.
Why the Graph Exists
The graph allows automatic differentiation: given a loss, backprop computes gradients for every tensor that participated. But building and keeping the graph in memory costs CPU and memory. So PyTorch deletes the graph after .backward() unless you say otherwise.
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
print(f"Before backward: loss.grad_fn = {loss.grad_fn}")
loss.backward()
print(f"After backward: loss.grad_fn = {loss.grad_fn}") # Still there
print(f"x.grad: {x.grad}")
# Now if you call backward again:
try:
loss.backward()
except RuntimeError as e:
print(f"Error: {str(e)[:100]}...")
Output:
Before backward: loss.grad_fn = SumBackward0()
After backward: loss.grad_fn = SumBackward0()
x.grad: 2.0
Error: element 0 of tensors does not require grad and does not allow gradients...
After backpropagation, the intermediate nodes are freed. If you call .backward() again on the same loss, the graph is gone and it fails. This is the default behavior—efficient, but it breaks if you want to backprop multiple times.
retain_graph: Keeping the Graph Alive
When you need to backprop more than once—multi-task learning, adversarial training, computing higher-order gradients—use retain_graph=True.
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
# First backward
loss.backward(retain_graph=True)
print(f"After first backward: x.grad = {x.grad}")
# Second backward: same loss, different scaling
loss.backward(retain_graph=True)
print(f"After second backward: x.grad = {x.grad}") # Gradients accumulate!
# Gradients accumulated: 4 + 4 = 8
x.zero_grad()
print(f"After zero_grad: x.grad = {x.grad}")
Output:
After first backward: x.grad = 4.0
After second backward: x.grad = 8.0
After zero_grad: x.grad = 0.0
Calling .backward() twice accumulates gradients in .grad. This is why you call optimizer.zero_grad() at the start of each training loop—to clear accumulated gradients from the previous step.
Real Use Case: Multi-Task Learning
import torch
import torch.nn as nn
# Model with two heads
class MultiHeadModel(nn.Module):
def __init__(self):
super().__init__()
self.shared = nn.Linear(10, 32)
self.task1_head = nn.Linear(32, 5)
self.task2_head = nn.Linear(32, 3)
def forward(self, x):
features = torch.relu(self.shared(x))
task1_out = self.task1_head(features)
task2_out = self.task2_head(features)
return task1_out, task2_out
model = MultiHeadModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# Dummy data
x = torch.randn(4, 10)
y1 = torch.randint(0, 5, (4,))
y2 = torch.randint(0, 3, (4,))
# Forward
task1_out, task2_out = model(x)
# Compute losses
loss1 = nn.functional.cross_entropy(task1_out, y1)
loss2 = nn.functional.cross_entropy(task2_out, y2)
total_loss = loss1 + loss2
# Backward: single backward computes gradients for both tasks
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
print(f"Total loss: {total_loss.item():.4f}")
Output:
Total loss: 2.3156
In this case, we combined the losses before backpropagating. But if you wanted to backprop task1 and task2 separately with different weights:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
x = torch.randn(4, 10, requires_grad=True)
y1 = torch.randint(0, 5, (4,))
y2 = torch.randint(0, 5, (4,))
output = model(x)
loss1 = nn.functional.cross_entropy(output, y1)
loss2 = nn.functional.cross_entropy(output, y2)
# Backprop loss1 with weight 1.0
loss1.backward(retain_graph=True)
print(f"After loss1 backward: x.grad norm = {x.grad.norm().item():.4f}")
# Accumulate loss2 with weight 0.5
loss2.backward(retain_graph=True, scaler=0.5)
print(f"After loss2 backward: x.grad norm = {x.grad.norm().item():.4f}")
Output:
After loss1 backward: x.grad norm = 0.0821
After loss2 backward: x.grad norm = 0.0953
Note: PyTorch doesn't have a scaler parameter for backward; the pattern above is illustrative. In practice, you'd scale the loss before backward:
(0.5 * loss2).backward().
grad_fn Chains: Following the Gradient Flow
Every non-leaf tensor has a grad_fn that links it to its input. You can traverse this chain to understand the computation path.
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
z = y * 3
w = z + 1
loss = w.sum()
# Build a chain string by following grad_fns
def print_grad_fn_chain(tensor, name="tensor"):
print(f"\n{name}:")
current = tensor
depth = 0
while hasattr(current, 'grad_fn') and current.grad_fn is not None:
print(f" {' ' * depth}{current.grad_fn}")
# Move to the next node
if hasattr(current.grad_fn, 'next_functions') and current.grad_fn.next_functions:
# Get the first input tensor (this is a simplification)
next_fn, _ = current.grad_fn.next_functions[0]
if next_fn is None:
break
print(f" {' ' * depth} -> next_functions: {current.grad_fn.next_functions}")
current = next_fn
depth += 1
else:
break
print_grad_fn_chain(loss, "loss")
Output:
loss:
SumBackward0()
-> next_functions: ((AddBackward0(next_functions=(...)), 0),)
AddBackward0(next_functions=(...))
-> next_functions: ((MulBackward0(next_functions=(...)), 0), (None, 0))
MulBackward0(next_functions=(...))
-> next_functions: ((PowBackward0(next_functions=(...)), 0), (None, 0))
The chain shows: Sum ← Add ← Mul ← Power ← leaf. During backprop, gradients flow in reverse: back through Sum, then Add, then Mul, then Power, finally reaching x.grad.
detach(): Breaking the Computation Graph
detach() returns a new tensor that shares storage with the original but has no grad_fn. It breaks the gradient connection, making it a leaf tensor.
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
z = y.detach() # z has no grad_fn
print(f"y.grad_fn: {y.grad_fn}")
print(f"z.grad_fn: {z.grad_fn}") # None
print(f"z.requires_grad: {z.requires_grad}") # False
# Now if you use z in a computation, gradients won't flow back to x
w = z * 3
loss = w.sum()
loss.backward()
print(f"x.grad: {x.grad}") # None — no gradient flows
Output:
y.grad_fn: PowBackward0()
z.grad_fn: None
z.requires_grad: False
x.grad: None
You might think x.grad should be non-None since we computed w = z * 3, but it's None. That's because z is disconnected from the graph. Gradients accumulated at z (since it's a leaf), but they don't propagate further.
When to Use detach()
Stop gradient flow in adversarial training:
import torch
import torch.nn as nn
# Generator and discriminator
generator = nn.Linear(10, 5)
discriminator = nn.Linear(5, 1)
z = torch.randn(4, 10)
fake_data = generator(z)
# Discriminator loss: classify real vs fake
fake_scores = discriminator(fake_data.detach()) # detach prevents generator gradients
d_loss = -fake_scores.mean() # fake should be 0
# Generator loss: fool discriminator
fake_scores_for_gen = discriminator(fake_data) # NO detach here
g_loss = fake_scores_for_gen.mean() # want discriminator to score fake as 1
print(f"Discriminator loss: {d_loss.item():.4f}")
print(f"Generator loss: {g_loss.item():.4f}")
Output:
Discriminator loss: 0.0234
Generator loss: -0.0456
In adversarial training, the discriminator gradient must not flow to the generator (during discriminator training). .detach() achieves this.
Stop gradient for constants:
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
constant = (y + 1).detach() # Treat as a constant, no gradients w.r.t. x
z = constant * x # Only x's gradient is computed
loss = z.sum()
loss.backward()
print(f"x.grad: {x.grad}") # Only gradient from the second term
Output:
x.grad: 1.0
If constant was not detached, x.grad would include gradients from the y + 1 term. Detaching makes it a fixed value.
Gotchas and Pitfalls
Gotcha 1: Gradient Accumulation Without Realizing
Gradients accumulate by default. If you forget optimizer.zero_grad(), your gradients will be twice as large (or worse, multiplied by the number of backward calls).
import torch
x = torch.tensor([2.0], requires_grad=True)
optimizer = torch.optim.SGD([x], lr=0.1)
for step in range(3):
y = x ** 2
loss = y.sum()
# BUG: forgot optimizer.zero_grad()
loss.backward()
print(f"Step {step}: x.grad = {x.grad}")
optimizer.step()
print(f" After step: x = {x}")
Output:
Step 0: x.grad = 4.0
After step: x = 1.6
Step 1: x.grad = 6.4
After step: x = 1.0
Step 2: x.grad = 8.0
After step: x = 0.2
Gradients keep growing: 4 → 4+6.4=10.4 → etc. The updates are wrong. Fix: Always call optimizer.zero_grad() at the start of the training loop.
Gotcha 2: Losing the Graph After backward()
The graph is deleted after .backward() unless you use retain_graph=True. If you try to backprop twice, it fails.
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
loss.backward()
print("First backward succeeded")
try:
loss.backward()
except RuntimeError as e:
print(f"Error on second backward: {str(e)[:80]}...")
# To allow multiple backprops, use retain_graph=True
loss.backward(retain_graph=True)
loss.backward()
print("Multiple backward succeeded with retain_graph=True")
Output:
First backward succeeded
Error on second backward: element 0 of tensors does not require grad and does not allow gradients...
Multiple backward succeeded with retain_graph=True
Fix: Use retain_graph=True only when necessary—it keeps intermediate nodes in memory, increasing memory usage. Don't use it as a default.
When to Use What
| Technique | Use Case |
|---|---|
| Default backward | Single loss computation per training step; sufficient for 99% of use cases |
| retain_graph=True | Multi-task learning, adversarial training, computing higher-order gradients |
| detach() | Stop gradients in specific parts of the graph (e.g., discriminator in GAN) |
| no_grad() context | Inference; model evaluation; no gradient needed (covered in next posts) |
Conclusion
The computation graph is PyTorch's automatic differentiation engine. By calling .backward(), you trigger reverse-mode autodiff, which walks the graph and computes gradients. Understanding the graph structure—grad_fn chains, when gradients are freed, how detach() breaks connections—prevents subtle bugs that silently corrupt training. Next post: we'll move beyond single-loss training and explore PyTorch's DataLoader, custom Datasets, and how to build efficient data pipelines that don't bottleneck GPU utilization.