Debugging PyTorch models

Shape mismatches, device and dtype errors, NaN losses, exploding gradients, autograd graph errors, and using hooks to inspect what the layers actually see.

Shape, device and dtype errors

# a shape trace that fits in a few lines and pays for itself immediately
def trace_shapes(model, sample, name="model"):
    hooks = []

    def make_hook(module_name):
        def hook(module, inputs, output):
            in_shape = tuple(inputs[0].shape) if inputs else None
            out_shape = tuple(output.shape) if isinstance(output, torch.Tensor) else type(output).__name__
            print(f"{module_name:45s} {str(in_shape):22s} -> {out_shape}")
        return hook

    for n, m in model.named_modules():
        if len(list(m.children())) == 0:
            hooks.append(m.register_forward_hook(make_hook(n)))

    model.eval()
    with torch.no_grad():
        model(sample.unsqueeze(0))

    for h in hooks:
        h.remove()

trace_shapes(model, val_images[0])

NaN, divergence and dead units

# narrow down where the NaN first appears
torch.autograd.set_detect_anomaly(True)     # slow, use only while debugging

def find_first_nan(model, x):
    activations = {}
    def hook(name):
        def fn(module, inputs, output):
            if isinstance(output, torch.Tensor) and not torch.isfinite(output).all():
                activations[name] = "non-finite output"
        return fn

    handles = [m.register_forward_hook(hook(n)) for n, m in model.named_modules()
               if len(list(m.children())) == 0]
    model(x)
    for h in handles:
        h.remove()
    return activations

print(find_first_nan(model, x_bad))

# a dead ReLU detector: a unit that outputs zero for every example is gone
def dead_relu_report(model, loader, threshold=0.0):
    counts, totals = {}, {}
    def hook(name):
        def fn(module, inputs, output):
            with torch.no_grad():
                alive = (output > threshold).any(dim=0).float().mean().item()
                counts[name] = counts.get(name, 0) + alive
                totals[name] = totals.get(name, 0) + 1
        return fn

    handles = [m.register_forward_hook(hook(n)) for n, m in model.named_modules()
               if isinstance(m, torch.nn.ReLU)]
    with torch.no_grad():
        for x, _ in loader:
            model(x)
    for h in handles:
        h.remove()
    return {n: counts[n] / totals[n] for n in counts}

• NaN arrives from three places: an exploding forward pass, an unstable loss (log of zero, division by zero), or an exploding gradient. Check in that order.
• Lower the learning rate by 10x as the first experiment. If the NaN disappears, it was divergence and the fix is numerical, not architectural.
• Dead ReLUs come from a learning rate that is too high or a bad initialisation. Switching to GELU or LeakyReLU often fixes it; scaling the initialisation fixes the cause.
• Wrap the loop with torch.autograd.set_detect_anomaly(True) only while hunting. It slows training substantially and should be removed afterwards.

Inspecting and modifying with hooks

import torch
import torch.nn as nn

# 1. a forward hook that captures an intermediate activation
captured = {}
def capture(name):
    def hook(module, inputs, output):
        captured[name] = output.detach()
    return hook

handle = model.layer4.register_forward_hook(capture("layer4"))
model(x)
print(captured["layer4"].shape, captured["layer4"].mean().item())
handle.remove()

# 2. a forward pre-hook that modifies the input to a layer
def clamp_inputs(module, inputs):
    return (inputs[0].clamp(-10, 10),)

model.fc.register_forward_pre_hook(clamp_inputs)

# 3. a full backward hook that reports gradient statistics
def grad_stats(module, grad_input, grad_output):
    g = grad_output[0]
    print(module.__class__.__name__, "grad mean/std:",
          round(g.mean().item(), 6), round(g.std().item(), 6))

model.fc.register_full_backward_hook(grad_stats)

# 4. gradient check for a hand-written layer
layer = CustomLayer()
x = torch.randn(4, 16, dtype=torch.double, requires_grad=True)
torch.autograd.gradcheck(layer.double(), (x,), eps=1e-6, atol=1e-4)

Hooks are also the safest way to probe a model you did not write: they attach from outside the code, capture exactly the tensor you want, and can be removed without editing the architecture.

FAQ

Related

Building networks: layers, containers and initialisation Exporting and deploying models

Last refreshed 2026-09-18.