Introduction

If you want the gradient of each layer's output in PyTorch, checking param.grad is not enough. Parameter gradients tell you how the loss changes with respect to weights and biases, while output gradients tell you what flowed backward through the activations produced by each layer.

The Key Distinction: Parameter Gradient Versus Output Gradient

After loss.backward(), PyTorch stores gradients for leaf parameters such as layer.weight.grad. That is useful for optimization, but it does not directly answer "what gradient hit this layer's output tensor?"

To inspect output gradients, you usually need one of these patterns:

• keep a reference to the intermediate tensor and call retain_grad()
• register hooks on module outputs during the forward pass

For per-layer debugging, output hooks are usually the most practical.

A Reliable Pattern with Forward Hooks and Tensor Hooks

One robust approach is:

1. register a forward hook on each layer
2. inside that hook, attach a tensor hook to the output
3. run a normal forward and backward pass
4. inspect the captured gradients

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5
6class Net(nn.Module):
7    def __init__(self):
8        super().__init__()
9        self.fc1 = nn.Linear(4, 8)
10        self.relu = nn.ReLU()
11        self.fc2 = nn.Linear(8, 1)
12
13    def forward(self, x):
14        x = self.fc1(x)
15        x = self.relu(x)
16        x = self.fc2(x)
17        return x
18
19
20model = Net()
21captured_output_grads = {}
22handles = []
23
24
25def make_hook(name):
26    def forward_hook(module, inputs, output):
27        if isinstance(output, torch.Tensor):
28            output.retain_grad()
29
30            def save_grad(grad):
31                captured_output_grads[name] = grad.detach().clone()
32
33            output.register_hook(save_grad)
34    return forward_hook
35
36
37for name, module in model.named_modules():
38    if name and len(list(module.children())) == 0:
39        handles.append(module.register_forward_hook(make_hook(name)))
40
41
42x = torch.randn(3, 4)
43y = torch.randn(3, 1)
44
45prediction = model(x)
46loss = F.mse_loss(prediction, y)
47loss.backward()
48
49for name, grad in captured_output_grads.items():
50    print(name, grad.shape, grad.norm().item())
51
52for handle in handles:
53    handle.remove()

This gives you the gradient that flowed through each captured layer output.

Why This Works Well

PyTorch's autograd system computes gradients on the backward pass only for tensors that matter to the loss. By attaching a tensor hook to the actual layer output, you inspect the activation gradient directly rather than inferring it indirectly from parameter gradients.

This is also more precise than a generic print statement inside the training loop because it lets you record gradients per layer, save them, compare norms, and inspect anomalies systematically.

Using register_full_backward_hook

PyTorch also provides register_full_backward_hook, which receives grad_input and grad_output for a module during backpropagation.

1def backward_hook(module, grad_input, grad_output):
2    if grad_output and grad_output[0] is not None:
3        print(module.__class__.__name__, grad_output[0].norm().item())
4
5handle = model.fc1.register_full_backward_hook(backward_hook)

This can be useful for quick debugging, but module backward hooks can be trickier than tensor hooks when modules return tuples, use in-place operations, or participate in more complex graphs. For many real debugging sessions, capturing output tensor hooks is the clearer option.

Checking a Single Saved Intermediate Tensor

If you only care about one specific activation, you can keep the tensor reference and call retain_grad() on it.

1x = torch.randn(2, 4)
2h = model.fc1(x)
3h.retain_grad()
4out = model.fc2(torch.relu(h))
5loss = out.sum()
6loss.backward()
7
8print(h.grad)

This is great for one-off inspection, but it does not scale nicely when you want every layer.

What to Look For

Once you capture output gradients, a few checks are especially useful:

• gradients that are all zeros, which may indicate dead activations or disconnected paths
• extremely tiny norms, which can suggest vanishing gradients
• very large norms, which may indicate exploding gradients
• unexpected None values, which often mean the tensor was not part of the backward path

Gradient inspection is most useful when paired with layer names, tensor shapes, and summary statistics such as min, max, and norm.

Common Pitfalls

The first pitfall is reading layer.weight.grad and assuming it represents the output gradient. It does not.

Another pitfall is forgetting retain_grad() on non-leaf tensors when you expect to inspect .grad directly.

A third pitfall is registering hooks and never removing them. That can create confusing repeated output or memory issues in longer-running training processes.

Finally, remember that if a tensor does not affect the chosen loss, it may not receive a gradient at all. In that case, hooks may not fire or gradients may be None.

Summary

• Output gradients are different from parameter gradients
• The most practical per-layer method is a forward hook that attaches a tensor hook to each output
• 'register_full_backward_hook can help, but tensor hooks are often easier to reason about'
• Use retain_grad() when inspecting a specific intermediate activation
• Watch for zero, tiny, huge, or missing gradients when debugging training behavior