From c65524d095936a0132976e8306b2f22b4f3ae7b4 Mon Sep 17 00:00:00 2001 From: Void Agent Date: Wed, 29 Jul 2026 19:45:41 +0100 Subject: jlens.py: fix hooks — don't detach activations, preserve gradient graph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/jlens.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) (limited to 'src') diff --git a/src/jlens.py b/src/jlens.py index 0cff809..a610e43 100644 --- a/src/jlens.py +++ b/src/jlens.py @@ -79,10 +79,8 @@ def compute_jlens_single_token(model, token_id, dataloader, layer_idx, device='c def make_hook(): def hook(module, input, output): - # output is (B, T, n_embd) - # Detach then require grad so we can compute gradient through it - activations['resid'] = output.detach().requires_grad_(True) - return activations['resid'] + # DON'T detach — need gradients to flow through + activations['resid'] = output return hook # Find the target layer @@ -178,8 +176,7 @@ def compute_jlens_all_layers(model, dataloader, n_layers, vocab_size, device='cu def make_hook(resid_dict): def hook(module, input, output): - resid_dict['val'] = output.detach().requires_grad_(True) - return resid_dict['val'] + resid_dict['val'] = output # Don't detach return hook target_block = model.transformer.h[layer_idx] -- cgit v1.2.3