jlens.py: fix hooks — don't detach activations, preserve gradient graph

1 files changed, 3 insertions, 6 deletions

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]