Read the middle of a model

You want the output of a layer in the middle of a model you did not write. The obvious options are bad: editing forward to return extra values changes the model’s interface for everyone, and reimplementing the forward pass up to that point duplicates the model.

A hook attaches to the layer instead.

captured = {}

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

handle = model.layer2.register_forward_hook(hook)
model(x)
handle.remove()

The hook fires after layer2 runs, receives the module, its inputs and its output, and is otherwise ordinary Python.

The two details that matter

detach in the hook. Without it, captured holds a tensor still attached to the graph, which keeps that whole forward pass alive for as long as you keep the dictionary. Doing this once per batch is a memory leak that grows until the run dies.

handle.remove(). Hooks persist on the module. Registering one per iteration and never removing gives you a list of hooks that grows every step, all of them running. remove() in a finally, or a context manager, is the reliable shape.

The variants

register_forward_pre_hook     # before, sees and can replace the inputs
register_forward_hook         # after, can replace the output by returning
register_full_backward_hook   # gradients flowing through the module

Returning a value from a forward hook replaces the output, which is how you patch a model’s behaviour without touching it. It is also how you break one by accident: a hook that ends with a stray expression returns it.

Where this is the standard tool

Feature extraction, attention-map visualisation, activation statistics, gradient-based attribution, and most interpretability tooling. nn.Module is designed for it, and it is the reason you should almost never need to fork a model to look inside it.

Your task

def capture_middle(x_values: list) -> dict

Build nn.Sequential(Linear, ReLU, Linear) with known weights, capture the output of the ReLU with a hook, run the model, remove the hook, and return the captured values, the final output, and the number of hooks left on the module afterwards.