Function hook

scripts/data/prepare_target_cache.py:97–98

(_module, _inputs, output)

    def capture_layer(layer_id: int):
        def hook(_module, _inputs, output):
            captured_hidden_states[layer_id] = _get_hook_tensor(output).detach()

        return hook