PyTorch Autograd Debugging: requires_grad, detach, backward() and NaN Gradients
PyTorch: Zero to Hero — Step 02
In the previous post we treated tensor shapes as a debugging problem rather than a mathematical vocabulary exercise.
We are going to do the same thing with autograd.
If you use PyTorch for any serious amount of time, you eventually see errors like:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
or:
RuntimeError: Trying to backward through the graph a second time...
or worse:
loss = nan
with no obvious exception at all.
This post is a programmer’s guide to answering five questions:
1. Is PyTorch recording this computation?
2. Which tensors are leaves?
3. Where should gradients appear?
4. Where did the graph get cut?
5. Where did NaN or Inf enter the backward pass?
If you can answer those questions quickly, most autograd bugs stop being mysterious.
The mental model: autograd records operations as you execute them
Start with the smallest useful example:
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
z = y ** 2
print(x)
print(y)
print(z)
Inspect the graph-related properties:
print("x.requires_grad:", x.requires_grad)
print("x.is_leaf:", x.is_leaf)
print("x.grad_fn:", x.grad_fn)
print("y.requires_grad:", y.requires_grad)
print("y.is_leaf:", y.is_leaf)
print("y.grad_fn:", y.grad_fn)
print("z.requires_grad:", z.requires_grad)
print("z.is_leaf:", z.is_leaf)
print("z.grad_fn:", z.grad_fn)
You should see the important distinction:
x: leaf tensor, requires grad, no grad_fn
y: non-leaf tensor, requires grad, has grad_fn
z: non-leaf tensor, requires grad, has grad_fn
Think about the forward computation as:
flowchart TD
X["x (leaf, requires_grad)"] --> mul["×3 (MulBackward)"]
mul --> Y["y"]
Y --> sq["² (PowBackward)"]
sq --> Z["z"]
Then:
z.backward()
walks backward through those recorded operations:
flowchart RL
Z --> sq
sq --> Y
Y --> mul
mul --> X
X --> |grad accumulated| XG["x.grad = 36"]
Now:
print(x.grad)
For:
z = (3x)^2 = 9x²
the derivative is:
dz/dx = 18x
At x = 2:
dz/dx = 36
So:
print(x.grad)
returns:
tensor([36.])
The first debugging rule is simple:
Inspect
requires_grad,is_leaf, andgrad_fnbefore guessing what autograd is doing.
A reusable tensor debugger
I use helpers like this constantly:
def debug_tensor(name, tensor):
print(
f"{name}: "
f"shape={tuple(tensor.shape)} "
f"dtype={tensor.dtype} "
f"device={tensor.device} "
f"requires_grad={tensor.requires_grad} "
f"is_leaf={tensor.is_leaf} "
f"grad_fn={type(tensor.grad_fn).__name__ if tensor.grad_fn else None} "
f"grad={'set' if tensor.grad is not None else None}"
)
Use it anywhere:
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
z = y.square()
debug_tensor("x", x)
debug_tensor("y", y)
debug_tensor("z", z)
This is far more useful than printing the tensor values alone.
requires_grad=True does not mean every tensor gets .grad
This surprises people.
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
z = y ** 2
z.backward()
print("x.grad:", x.grad)
print("y.grad:", y.grad)
print("z.grad:", z.grad)
x.grad is populated because x is a leaf tensor that requires gradients.
But y.grad is normally None.
Why?
Because PyTorch accumulates gradients into .grad for leaf tensors by default.
If you want the gradient of an intermediate value for debugging, explicitly retain it:
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
y.retain_grad()
z = y ** 2
z.backward()
print("x.grad:", x.grad)
print("y.grad:", y.grad)
Now y.grad is available.
This is extremely useful when debugging a deep model:
hidden = layer1(x)
hidden.retain_grad()
output = layer2(hidden)
loss = criterion(output, target)
loss.backward()
print(hidden.grad)
If gradients are healthy before layer2 but broken after it, you just narrowed your search dramatically.
Error 1: element 0 of tensors does not require grad
The classic error:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Usually means the tensor on which you called backward() is not connected to an autograd graph.
Minimal reproduction
import torch
x = torch.tensor([2.0])
y = x * 3
loss = y.square().sum()
loss.backward()
This fails because:
print(x.requires_grad)
print(y.requires_grad)
print(loss.requires_grad)
all report False.
Fix
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
loss = y.square().sum()
loss.backward()
print(x.grad)
But in real code the bug is often subtler.
The graph was cut with detach()
Consider:
x = torch.tensor([2.0], requires_grad=True)
hidden = x * 3
hidden = hidden.detach()
loss = hidden.square().sum()
loss.backward()
Inspect it:
debug_tensor("x", x)
debug_tensor("hidden", hidden)
debug_tensor("loss", loss)
detach() deliberately returns a tensor disconnected from the current graph.
A common production bug looks more like this:
features = encoder(images)
features = features.detach()
logits = classifier(features)
loss = criterion(logits, labels)
loss.backward()
If you intended to train only the classifier, this can be correct.
If you intended to train the encoder too, you just silently froze the entire path before the classifier.
Debugging rule
When gradients unexpectedly stop:
# conceptually search your codebase for:
detach
no_grad
inference_mode
requires_grad_(False)
item()
numpy()
Those are common graph boundaries.
.item() is a graph exit
This is fine:
loss_value = loss.item()
print(loss_value)
because you are logging a scalar.
This is not fine if you try to feed it back into the loss calculation:
part_a = model_a(x).mean()
part_b = model_b(x).mean()
bad_loss = torch.tensor(part_a.item() + part_b.item(), requires_grad=True)
bad_loss.backward()
You created a brand-new tensor from Python numbers.
The original graph is gone.
The correct version is simply:
loss = part_a + part_b
loss.backward()
You can log afterward:
print(loss.item())
A good rule:
Use
.item()for observation, not computation.
NumPy can cut the graph too
This is another easy trap:
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
array = y.detach().cpu().numpy()
That is fine if you need NumPy for logging or visualization.
But if you compute in NumPy and then convert back:
import numpy as np
array = y.detach().cpu().numpy()
array = np.square(array)
z = torch.tensor(array, requires_grad=True)
z.sum().backward()
z belongs to a new graph.
There is no path back to x.
print(x.grad)
will still be None.
If the operation exists in PyTorch, keep it in PyTorch:
z = y.square()
z.sum().backward()
Error 2: gradients keep getting larger every iteration
PyTorch accumulates gradients.
That means this:
x = torch.tensor([2.0], requires_grad=True)
for step in range(3):
y = x ** 2
y.backward()
print(step, x.grad)
prints increasing gradients.
For x = 2, each backward pass contributes 4:
step 0 tensor([4.])
step 1 tensor([8.])
step 2 tensor([12.])
This is intentional.
Gradient accumulation is useful.
But most training loops want to reset gradients each iteration.
Manual tensors
for step in range(3):
if x.grad is not None:
x.grad.zero_()
y = x ** 2
y.backward()
print(step, x.grad)
Real model
Typical code:
for inputs, targets in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
A common alternative is:
optimizer.zero_grad(set_to_none=True)
Then parameters without gradients remain None rather than being filled with zero tensors.
This can also be useful while debugging:
for name, param in model.named_parameters():
print(name, param.grad is None)
Now None can tell you that a parameter did not participate in the backward path.
You can also visualise the accumulation:
import matplotlib.pyplot as plt
x = torch.tensor([2.0], requires_grad=True)
acc = []
for _ in range(10):
y = x ** 2
y.backward()
acc.append(x.grad.item())
plt.plot(range(1,11), acc, marker='o')
plt.xlabel('Backward call')
plt.ylabel('x.grad')
plt.title('Gradient accumulation without zeroing')
plt.grid(True)
plt.show()
A useful gradient report
Add this to your debugging toolkit:
def gradient_report(model):
for name, param in model.named_parameters():
if not param.requires_grad:
status = "FROZEN"
elif param.grad is None:
status = "NO_GRAD"
else:
grad = param.grad.detach()
status = (
f"mean={grad.abs().mean().item():.3e} "
f"max={grad.abs().max().item():.3e} "
f"finite={torch.isfinite(grad).all().item()}"
)
print(f"{name:50s} {status}")
Call it immediately after backward:
loss.backward()
gradient_report(model)
optimizer.step()
This answers several important questions at once:
Is the parameter trainable?
Did it receive a gradient?
Are the gradients tiny?
Are they enormous?
Are they finite?
Error 3: Trying to backward through the graph a second time
Minimal example:
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
loss.backward()
loss.backward()
The second call fails.
PyTorch normally frees saved intermediate values used by backward once they are no longer needed.
If you genuinely need multiple backward passes through the same graph:
loss.backward(retain_graph=True)
loss.backward()
But be careful.
retain_graph=True is often used as a plaster over a design bug.
Before adding it, ask:
Why am I reusing this graph?
Should I recompute the forward pass instead?
Am I accidentally carrying tensors from one training iteration into another?
Most ordinary training loops should rebuild the graph each forward pass:
for batch in dataloader:
optimizer.zero_grad()
output = model(batch)
loss = compute_loss(output)
loss.backward()
optimizer.step()
New forward pass, new graph.
Hidden-state bugs in recurrent code
Graph reuse problems commonly appear in loops where state is carried forward:
hidden = torch.zeros(hidden_size, requires_grad=True)
for x in sequence:
hidden = recurrent_step(hidden, x)
loss = compute_loss(hidden)
loss.backward()
Depending on what you intend, the graph can keep linking backward across iterations.
For truncated backpropagation through time you may deliberately detach state:
hidden = hidden.detach()
or:
hidden = hidden.detach().requires_grad_()
The key is intent.
detach() is neither good nor bad.
It is a graph boundary.
You should know why it is there.
Error 4: in-place operations break backward
An in-place operation modifies an existing tensor.
Many PyTorch methods signal this with a trailing underscore:
x.add_(1)
x.relu_()
x.zero_()
In-place operations can conflict with autograd because backward may need an earlier version of a tensor.
Example:
x = torch.tensor([2.0], requires_grad=True)
y = x * 3
z = y ** 2
y.add_(1)
z.backward()
Depending on the exact operation and tensors saved for backward, PyTorch may report that a variable needed for gradient computation was modified in-place.
The safer default during model development is:
y = y + 1
rather than:
y += 1
when that tensor participates in autograd and you do not specifically need the in-place behavior.
Do not mechanically ban every in-place operation.
But if backward starts complaining about version counters or modified variables, search for:
+=
-=
*=
/=
*_()
add_()
relu_()
copy_()
near the failing path.
Error 5: model.eval() did not disable gradients
This one is common because the names invite confusion.
model.eval()
output = model(x)
model.eval() changes the behavior of modules that have train/eval behavior, such as dropout and batch normalization.
It does not globally disable autograd.
If you want evaluation without gradient recording:
model.eval()
with torch.no_grad():
output = model(x)
or, for pure inference where its restrictions are acceptable:
model.eval()
with torch.inference_mode():
output = model(x)
These are separate concepts:
model.train() / model.eval()
controls module training behaviour
requires_grad
controls whether specific tensors/parameters participate in autograd
no_grad()
prevents operations in the block from being recorded for backward
inference_mode()
more aggressive inference-oriented autograd disabling
Do not treat them as interchangeable.
Debug the current gradient mode
If you suspect some surrounding helper disabled gradients:
print(torch.is_grad_enabled())
Example:
x = torch.tensor([2.0], requires_grad=True)
print("outside:", torch.is_grad_enabled())
with torch.no_grad():
print("inside:", torch.is_grad_enabled())
y = x * 3
print("y.requires_grad:", y.requires_grad)
You can temporarily re-enable gradients inside a disabled region:
with torch.no_grad():
a = x * 2
with torch.enable_grad():
b = x * 3
Inspect both:
print(a.requires_grad)
print(b.requires_grad)
When debugging framework code, callbacks, evaluation wrappers, or mixed training/inference pipelines, torch.is_grad_enabled() is one of the first things I would print.
Error 6: frozen parameters you forgot were frozen
Fine-tuning code often contains:
for param in model.encoder.parameters():
param.requires_grad = False
Later you change the experiment and expect the encoder to train.
It does not.
Create a report:
def parameter_report(model):
total = 0
trainable = 0
for name, param in model.named_parameters():
count = param.numel()
total += count
if param.requires_grad:
trainable += count
print(
f"{name:50s} "
f"shape={str(tuple(param.shape)):20s} "
f"requires_grad={param.requires_grad}"
)
print()
print(f"total parameters: {total:,}")
print(f"trainable parameters: {trainable:,}")
Run it before training:
parameter_report(model)
This catches a surprising number of mistakes.
You can also freeze and unfreeze modules explicitly:
model.encoder.requires_grad_(False)
model.classifier.requires_grad_(True)
and later:
model.encoder.requires_grad_(True)
Error 7: the optimizer does not know about your parameter
A parameter can receive a gradient and still never update.
Example:
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.layer1 = torch.nn.Linear(10, 10)
self.layer2 = torch.nn.Linear(10, 1)
def forward(self, x):
return self.layer2(torch.relu(self.layer1(x)))
Suppose you create the optimizer incorrectly:
optimizer = torch.optim.Adam(model.layer2.parameters(), lr=1e-3)
layer1 can participate in backward and have gradients, but the optimizer will not update it.
A quick check:
optimizer_parameter_ids = {
id(param)
for group in optimizer.param_groups
for param in group["params"]
}
for name, param in model.named_parameters():
print(
name,
"requires_grad=", param.requires_grad,
"has_grad=", param.grad is not None,
"in_optimizer=", id(param) in optimizer_parameter_ids,
)
This distinguishes:
NO_GRAD
from:
HAS_GRAD_BUT_NOT_OPTIMIZED
Those are completely different bugs.
Error 8: loss becomes NaN
This is where debugging gets interesting.
A training loop can run perfectly for hundreds of steps and then:
step=742 loss=nan
Do not immediately change the learning rate and hope.
Instrument the computation.
First: check the forward values
Reusable helper:
def assert_finite(name, tensor):
if not torch.isfinite(tensor).all():
bad = tensor[~torch.isfinite(tensor)]
raise RuntimeError(
f"{name} contains non-finite values: "
f"count={bad.numel()} "
f"sample={bad.flatten()[:10]}"
)
Use it aggressively:
hidden = layer1(x)
assert_finite("layer1 output", hidden)
hidden = torch.relu(hidden)
assert_finite("relu output", hidden)
logits = layer2(hidden)
assert_finite("logits", logits)
loss = criterion(logits, target)
assert_finite("loss", loss)
The first failure is usually more useful than the final NaN loss.
Second: inspect gradients immediately after backward
loss.backward()
for name, param in model.named_parameters():
if param.grad is None:
continue
if not torch.isfinite(param.grad).all():
raise RuntimeError(f"non-finite gradient in {name}")
Or include magnitude:
for name, param in model.named_parameters():
if param.grad is None:
continue
grad = param.grad.detach()
print(
name,
"mean=", grad.abs().mean().item(),
"max=", grad.abs().max().item(),
"finite=", torch.isfinite(grad).all().item(),
)
The difference between:
1e-6
and:
1e+23
matters.
Third: anomaly detection
PyTorch has a debugging mode specifically for autograd anomalies:
with torch.autograd.detect_anomaly():
output = model(x)
loss = criterion(output, target)
loss.backward()
or:
torch.autograd.set_detect_anomaly(True)
During debugging, this can identify the forward operation associated with a failing backward function and can raise when backward generates NaNs.
Do not leave it enabled in normal training.
It adds overhead.
Use it like a debugger, not a runtime feature.
A deliberately broken NaN example
import torch
x = torch.tensor([1.0, 0.0], requires_grad=True)
y = x / x
loss = y.sum()
print("y:", y)
print("loss:", loss)
loss.backward()
print("x.grad:", x.grad)
At 0 / 0, you get a NaN in the forward pass.
This is easy to catch with:
assert_finite("y", y)
The harder cases are those where the forward values still look plausible but backward becomes unstable.
That is exactly where anomaly detection and gradient hooks become useful.
Gradient hooks: inspect the backward pass while it happens
You can attach a hook to a tensor:
x = torch.tensor([2.0], requires_grad=True)
x.register_hook(lambda grad: print("gradient into x:", grad))
y = x ** 2
y.backward()
For model parameters:
handles = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
handle = param.register_hook(
lambda grad, name=name: print(
name,
"mean=", grad.abs().mean().item(),
"max=", grad.abs().max().item(),
)
)
handles.append(handle)
Run a backward pass:
loss.backward()
Then remove the hooks:
for handle in handles:
handle.remove()
For large models, printing every parameter is too noisy.
Make the hook conditional:
def make_grad_hook(name):
def hook(grad):
if not torch.isfinite(grad).all():
print(f"NON-FINITE GRADIENT: {name}")
elif grad.abs().max() > 1000:
print(
f"LARGE GRADIENT: {name} "
f"max={grad.abs().max().item():.3e}"
)
return hook
handles = [
param.register_hook(make_grad_hook(name))
for name, param in model.named_parameters()
if param.requires_grad
]
Now the instrumentation only screams when something interesting happens.
Gradient clipping: useful, but not a diagnosis
Exploding gradients are sometimes controlled with clipping:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
Typical placement:
optimizer.zero_grad(set_to_none=True)
output = model(inputs)
loss = criterion(output, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
But clipping can hide the symptom without explaining the cause.
If your gradient norm suddenly jumps from:
2.4
to:
8.7e18
you still want to know why.
Use clipping as a training technique when appropriate.
Do not use it as a substitute for instrumentation.
Measure the total gradient norm
def grad_norm(model, norm_type=2.0):
grads = [
param.grad.detach()
for param in model.parameters()
if param.grad is not None
]
if not grads:
return torch.tensor(0.0)
norms = torch.stack([
torch.linalg.vector_norm(g, ord=norm_type)
for g in grads
])
return torch.linalg.vector_norm(norms, ord=norm_type)
Use it:
loss.backward()
print("grad norm:", grad_norm(model).item())
optimizer.step()
Log that over time.
A gradient norm curve can expose instability before the loss visibly explodes.
Why masking after an invalid operation may not save the gradient
This is a subtle one.
Suppose you divide before masking:
x = torch.tensor([1.0, 1.0], requires_grad=True)
divisor = torch.tensor([0.0, 1.0])
ratio = x / divisor
masked = ratio[divisor != 0]
loss = masked.sum()
print(ratio)
print(masked)
You might reason:
The invalid result was removed before the loss.
But autograd recorded the division operation in the graph.
The better design is to avoid constructing invalid operations in the first place.
For example:
mask = divisor != 0
safe_ratio = x[mask] / divisor[mask]
loss = safe_ratio.sum()
This principle generalizes:
Do not compute invalid values and hope masking later makes the backward graph safe.
Prevent the invalid operation when possible.
backward() on non-scalar outputs
This catches people moving from losses to lower-level autograd experiments.
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
y.backward()
A vector output does not implicitly define a single scalar gradient seed.
The simplest solution is often to reduce it:
loss = y.sum()
loss.backward()
Now:
print(x.grad)
returns:
tensor([2., 4., 6.])
Or explicitly provide the vector used in the vector-Jacobian product:
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
y.backward(torch.ones_like(y))
print(x.grad)
For training code, losses are normally scalars after reduction.
If yours is not, inspect its shape:
print(loss.shape)
before calling backward.
torch.autograd.grad() when you want gradients returned, not accumulated
Sometimes .backward() is not the right debugging tool.
Use:
x = torch.tensor([2.0], requires_grad=True)
y = x ** 3
(dx,) = torch.autograd.grad(y, x)
print(dx)
Output:
tensor([12.])
Unlike normal .backward() usage, torch.autograd.grad() returns the requested gradients directly.
This is useful for:
custom losses
gradient penalties
meta-learning
higher-order derivatives
unit-testing derivative behaviour
small debugging probes
Example gradient test:
def f(x):
return x ** 3 + 2 * x
x = torch.tensor([4.0], requires_grad=True)
y = f(x)
(dx,) = torch.autograd.grad(y, x)
expected = 3 * x.detach() ** 2 + 2
print("autograd:", dx)
print("expected:", expected)
assert torch.allclose(dx, expected)
Autograd is software.
You can test it like software.
Finite differences: test your gradient numerically
If you suspect a custom computation, compare autograd with a numerical derivative.
import torch
def f(x):
return (x ** 3 + 2 * x).sum()
x = torch.tensor([4.0], requires_grad=True)
y = f(x)
y.backward()
autograd_gradient = x.grad.detach().clone()
with torch.no_grad():
epsilon = 1e-4
x_plus = x.detach() + epsilon
x_minus = x.detach() - epsilon
numerical_gradient = (
f(x_plus) - f(x_minus)
) / (2 * epsilon)
print("autograd:", autograd_gradient)
print("numerical:", numerical_gradient)
For complicated custom math, this gives you an independent check.
PyTorch also provides gradient-checking utilities for custom autograd functions, but the finite-difference idea is worth understanding directly.
A minimal model-debugging harness
Here is a compact harness worth keeping around:
import torch
def assert_finite(name, tensor):
if not torch.isfinite(tensor).all():
raise RuntimeError(f"non-finite tensor: {name}")
def report_parameters(model):
for name, param in model.named_parameters():
print(
f"{name:40s} "
f"shape={str(tuple(param.shape)):18s} "
f"requires_grad={str(param.requires_grad):5s}"
)
def report_gradients(model):
for name, param in model.named_parameters():
if param.grad is None:
print(f"{name:40s} grad=None")
continue
grad = param.grad.detach()
print(
f"{name:40s} "
f"mean={grad.abs().mean().item():.3e} "
f"max={grad.abs().max().item():.3e} "
f"finite={torch.isfinite(grad).all().item()}"
)
def debug_training_step(model, optimizer, criterion, inputs, targets):
model.train()
optimizer.zero_grad(set_to_none=True)
assert_finite("inputs", inputs)
outputs = model(inputs)
assert_finite("outputs", outputs)
loss = criterion(outputs, targets)
assert_finite("loss", loss)
loss.backward()
report_gradients(model)
optimizer.step()
return loss.detach()
Use it before debugging a 2,000-line training pipeline.
Reduce the problem to one batch and one step.
loss = debug_training_step(
model,
optimizer,
criterion,
inputs,
targets,
)
print("loss:", loss.item())
If one step is broken, distributed training will not fix it.
Mixed precision will not fix it.
A bigger GPU will not fix it.
Fix the smallest reproducible training step first.
Detect parameters that never change
Sometimes gradients look correct but the model still does not learn.
Snapshot parameters before the step:
before = {
name: param.detach().clone()
for name, param in model.named_parameters()
}
Run the step:
optimizer.zero_grad(set_to_none=True)
output = model(inputs)
loss = criterion(output, targets)
loss.backward()
optimizer.step()
Compare:
for name, param in model.named_parameters():
changed = not torch.equal(before[name], param.detach())
print(name, "changed=", changed)
For more detail:
for name, param in model.named_parameters():
delta = (param.detach() - before[name]).abs()
print(
name,
"mean_delta=", delta.mean().item(),
"max_delta=", delta.max().item(),
)
This separates three stages that programmers often collapse into one idea:
flowchart TD
A[Forward pass] --> B[Gradient computation]
B --> C[Optimizer update]
style A fill:#e1f5fe
style B fill:#fff9c4
style C fill:#c8e6c9
A parameter can fail at any one of those stages.
Detect dead gradient paths automatically
Create a function that labels every parameter:
def classify_parameter_gradients(model, optimizer=None):
optimizer_ids = None
if optimizer is not None:
optimizer_ids = {
id(param)
for group in optimizer.param_groups
for param in group["params"]
}
rows = []
for name, param in model.named_parameters():
if not param.requires_grad:
state = "FROZEN"
elif param.grad is None:
state = "NO_GRAD"
elif not torch.isfinite(param.grad).all():
state = "NON_FINITE_GRAD"
elif optimizer_ids is not None and id(param) not in optimizer_ids:
state = "NOT_IN_OPTIMIZER"
elif param.grad.abs().max() == 0:
state = "ZERO_GRAD"
else:
state = "OK"
rows.append((name, state))
return rows
Then:
for name, state in classify_parameter_gradients(model, optimizer):
print(f"{name:50s} {state}")
We can see the classification logic as a decision tree:
flowchart TD
A[Parameter] --> B{requires_grad?}
B -- No --> C[FROZEN]
B -- Yes --> D{grad is None?}
D -- Yes --> E[NO_GRAD]
D -- No --> F{finite grad?}
F -- No --> G[NON_FINITE_GRAD]
F -- Yes --> H{in optimizer?}
H -- No --> I[NOT_IN_OPTIMIZER]
H -- Yes --> J{max grad == 0?}
J -- Yes --> K[ZERO_GRAD]
J -- No --> L[OK]
That is the beginning of a real training diagnostic tool.
A programmer’s autograd debugging checklist
When training breaks, work through this in order.
flowchart TD
S[Start debugging] --> Q1[Is grad mode enabled?]
Q1 -->|No| FIX1[Enable or exit no_grad]
Q1 -->|Yes| Q2[Does loss require_grad?]
Q2 -->|No| FIX2[Find broken graph chain]
Q2 -->|Yes| Q3[Are all intended params trainable?]
Q3 -->|No| FIX3[Set requires_grad=True]
Q3 -->|Yes| Q4[Do params have grad?]
Q4 -->|Some None| FIX4[Find where graph is detached]
Q4 -->|All have grad| Q5[Are grads finite?]
Q5 -->|No| FIX5[NaN/inf source, check forward]
Q5 -->|Yes| Q6[Gradient magnitudes sane?]
Q6 -->|No| FIX6[Exploding/vanishing, check norms]
Q6 -->|Yes| Q7[Params in optimizer?]
Q7 -->|No| FIX7[Add to param_groups]
Q7 -->|Yes| Q8[Do params change after step?]
Q8 -->|No| FIX8[Update step not effective, check LR]
Q8 -->|Yes| DONE[Training step OK]
1. Is grad mode enabled?
print(torch.is_grad_enabled())
2. Does the loss require grad?
print(loss.requires_grad)
print(loss.grad_fn)
If not, walk backward through the tensors that created it.
3. Are the expected parameters trainable?
for name, param in model.named_parameters():
print(name, param.requires_grad)
4. Did every expected parameter receive a gradient?
After:
loss.backward()
inspect:
for name, param in model.named_parameters():
print(name, param.grad is None)
5. Are gradients finite?
for name, param in model.named_parameters():
if param.grad is not None:
print(name, torch.isfinite(param.grad).all().item())
6. Are gradient magnitudes sane?
for name, param in model.named_parameters():
if param.grad is not None:
print(name, param.grad.abs().max().item())
7. Is the parameter actually in the optimizer?
optimizer_parameter_ids = {
id(param)
for group in optimizer.param_groups
for param in group["params"]
}
for name, param in model.named_parameters():
print(name, id(param) in optimizer_parameter_ids)
8. Did the parameter change after optimizer.step()?
Snapshot and compare.
9. Search for graph boundaries
detach()
item()
numpy()
no_grad()
inference_mode()
requires_grad_(False)
10. Search for in-place modifications
+=
-=
*=
/=
method_name_()
11. Turn on anomaly detection
with torch.autograd.detect_anomaly():
loss.backward()
12. Reduce to one batch
Remove:
distributed training
mixed precision
gradient accumulation
callbacks
logging frameworks
schedulers
augmentation
compilation
until the smallest broken training step remains.
Then add complexity back.
The debugging loop I want you to internalize
When a model is not learning, do not start with:
Maybe Adam is bad.
Maybe I need a different architecture.
Maybe I need more data.
Start with execution evidence:
flowchart TD
A[Does forward pass contain finite values?] -->|Yes| B[Does loss belong to autograd graph?]
A -->|No| Aerr[Fix forward]
B -->|Yes| C[Do intended params require grad?]
B -->|No| Berr[Fix graph detachment]
C -->|Yes| D[Did those params receive grad?]
C -->|No| Cerr[Set requires_grad=True]
D -->|Yes| E[Are grads finite and non-zero?]
D -->|No| Derr[Check graph path]
E -->|Yes| F[Does optimizer contain those params?]
E -->|No| Eerr[Fix NaN/zero grad]
F -->|Yes| G["Did optimizer.step() change them?"]
F -->|No| Ferr[Add to optimizer]
G -->|Yes| Success[Learning pipeline is functional]
G -->|No| Gerr[Check LR, update rule]
That sequence turns “the model isn’t learning” into a normal software-debugging problem.
And that is exactly how programmers should approach PyTorch.
Challenge: deliberately break autograd
Create this tiny model:
import torch
import torch.nn as nn
class TinyNet(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(4, 8)
self.fc2 = nn.Linear(8, 1)
def forward(self, x):
x = self.fc1(x)
x = torch.relu(x)
return self.fc2(x)
model = TinyNet()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
x = torch.randn(16, 4)
target = torch.randn(16, 1)
Then deliberately introduce these bugs one at a time:
Bug A
output = model(x).detach()
loss = ((output - target) ** 2).mean()
loss.backward()
Bug B
for param in model.fc1.parameters():
param.requires_grad = False
Bug C
Create the optimizer with only:
model.fc2.parameters()
Bug D
Remove:
optimizer.zero_grad()
from the loop.
Bug E
Introduce a dangerous division:
output = output / 0
For each bug, use the tools in this post to identify exactly which stage failed.
Do not fix the bug first.
Prove where it is.
That is the skill we are building.
Where the series goes next
Step 00 — What Are We Actually Doing?
Step 01 — Tensor Shapes, Broadcasting and Shape Errors
Step 02 — Autograd Debugging: requires_grad, detach, backward() and NaNs
Step 03 — Build a Neural Network Without nn.Module
Step 04 — Now Let PyTorch Do the Plumbing
Step 05 — Real Data: Dataset, DataLoader and Training Loops
Step 06 — CNNs: Teaching PyTorch to See
Step 07 — Attention and Transformers From Scratch
Step 08 — Train Something Real
Step 09 — Performance, Compilation and Scale
Step 10 — Build a Small Language Model From Scratch
The next post will deliberately remove most of PyTorch’s neural-network abstractions.
We will build a trainable neural network using tensors, matrix multiplication, activation functions and autograd — but without nn.Module.
By that point, nn.Module will stop looking like the thing that creates neural networks.
It will look like what it actually is:
useful plumbing around computations we already understand.