3. Capture Intermediate Activations with a Forward Hook

Problem

Interpretability work needs to inspect what a specific layer actually computed on a given input, without rewriting the model's forward() method to return intermediate values. PyTorch's forward hooks let you attach a callback to any submodule that fires right after that submodule's forward() runs, handing you its output before it's discarded.

This is the standard, non-invasive way to pull activations out of an off-the-shelf model for analysis, and it must clean up after itself — a hook left registered fires on every future forward pass and silently leaks memory.

Source: src/3_activation_capture_hook.py

def capture_activation(model: torch.nn.Module, layer_name: str, input_tensor: torch.Tensor) -> torch.Tensor: ...

>>> model = nn.Sequential(nn.Linear(4, 6), nn.ReLU())
>>> x = torch.randn(2, 4)
>>> capture_activation(model, "0", x).shape
torch.Size([2, 6])

Step-by-Step Approach

1. Resolve the target submodule from its dotted path with model.get_submodule(layer_name).
2. Define a hook function with signature (module, inputs, output) that stashes output.detach() somewhere the outer function can read afterward (a small mutable container like a dict or list works, since the hook closes over it).
3. Register the hook with submodule.register_forward_hook(hook), which returns a handle.
4. Run the full model forward pass on input_tensor — the hook fires automatically the moment the target submodule finishes.
5. Remove the hook with handle.remove() inside a try/finally, so it's cleaned up even if the forward pass raises.
6. Return the captured activation tensor.

The key insight is that a forward hook is just a callback wired into the module's own __call__ machinery — it doesn't need any changes to the model's source code, which is exactly why it's the standard tool for probing third-party or frozen models you don't want to modify.

Reference solution

import torch
import torch.nn as nn


def capture_activation(model: nn.Module, layer_name: str, input_tensor: torch.Tensor) -> torch.Tensor:
    submodule = model.get_submodule(layer_name)
    captured = {}

    def hook(module, inputs, output):
        captured["activation"] = output.detach()

    handle = submodule.register_forward_hook(hook)
    try:
        model(input_tensor)
    finally:
        # always remove the hook, even if the forward pass raises --
        # a leaked hook fires on every future forward call and leaks memory
        handle.remove()
    return captured["activation"]

Key Functions & Tricks

• model.get_submodule(layer_name) — resolves a dotted attribute path (e.g. "blocks.1.mlp") to the actual submodule object, including numeric indices into nn.Sequential.
• module.register_forward_hook(fn) — attaches fn(module, inputs, output) to fire right after that module's forward() returns; returns a removable handle.
• handle.remove() — detaches the hook; always call this, ideally in a finally block, to avoid leaking hooks across repeated calls.
• output.detach() — captured activations should not keep the whole forward graph alive in memory just for inspection.
• A closure-captured mutable container (dict/list) — the standard way to get a value out of a hook callback, since the hook's return value is ignored for a plain forward hook.

How to Recognize This Pattern