How do you capture a PyTorch model's intermediate activations without editing its forward()?

## The problem hooks solve You need the activations of layer 3 of a pretrained encoder — for a probing experiment, a feature extractor, a similarity search index, or just to see where a signal dies. The obvious approaches are bad: forking the library to add a `return_intermediates=True` flag couples you to that fork forever, and rebuilding the model as a `nn.Sequential` of its children silently drops anything the original `forward` did between modules (reshapes, residual adds, functional activations). Hooks are PyTorch's supported answer. They are callbacks the module dispatches during `__call__`, so they see real traffic through the real `forward`. ## The API `handle = module.register_forward_hook(fn)` where `fn(module, args, output)`: - `module` is the module the hook is attached to; - `args` is the tuple of positional inputs it was called with; - `output` is whatever `forward` returned. Returning `None` (the usual case) leaves the output untouched; returning a value **replaces** the module's output for the rest of the computation, which is how people patch activations for interpretability experiments. `register_forward_pre_hook(fn)` runs before `forward`, with `fn(module, args)`, and can replace the inputs. `register_full_backward_hook` exists for the gradient side. Both registration calls return a `RemovableHandle` whose `.remove()` deregisters the hook. That handle is the piece people forget. ## Finding the submodule Two reliable routes: - `model.get_submodule("encoder.layers.3.mlp")` — dotted-path lookup, raises `AttributeError` on a bad path, which is exactly what you want rather than silently hooking nothing. - iterate `model.named_modules()` and match on the name or on `isinstance`. This is how you hook *every* `nn.Conv2d` in one loop for a layer-wise statistics dump. The names are the same attribute paths that form `state_dict` keys, so anything you can see in a checkpoint you can address here. ## The two things that go wrong in production **Holding the graph.** If your hook does `store[name] = output`, you keep a tensor that still carries `grad_fn`, and through it the entire autograd graph for that iteration. Backward cannot free what you are still referencing. Memory climbs every step until the process OOMs — and because the growth is gradual it reads exactly like a slow leak elsewhere. Store `output.detach()`, and `.cpu()` too if you are collecting across many steps. **Leaked handles.** Hooks registered in a helper and never removed keep firing through validation, through export tracing, through the next experiment in the same process. A hook that appends to a list becomes an unbounded accumulator. The disciplined shape is a try/finally, or a small context manager that removes every handle it registered on exit. ## The silent no-op Hooks fire from `nn.Module.__call__`. If any code path invokes `submodule.forward(x)` directly, the hook does not run. You get an empty dictionary and no error at all. The same applies if you attached the hook to the wrong instance — for example to a freshly constructed layer rather than the one inside the model, or to `ddp_model.module`'s child while calling a differently wrapped object. When a hook seems dead, verify the identity of the module you hooked and confirm the call path goes through `model(x)`. ## Alternatives and when they are better - **`torch.fx` / `create_feature_extractor`-style graph rewriting** gives you a new module whose forward returns the nodes you asked for. It is more robust for a fixed extraction you will run many times, and it survives serialization; it is heavier to set up and can fail on models with data-dependent control flow. - **Subclassing and overriding `forward`** is fine when you own the model. Hooks earn their keep precisely when you do not. - **Export-time capture** is a different problem: a hook is a Python callback, so it does not survive `torch.export` or a traced graph. If your goal is to ship a model that returns intermediates, put that in the model, not in a hook. ## The shape of good hook code Register, run, remove, in a bounded scope. Detach what you keep. Address submodules by their dotted path so the code fails loudly when the architecture changes underneath you. Never leave a hook installed across phases of a job — activations captured during validation because a training hook was never removed is a genuinely painful bug to find.