diff options
| -rw-r--r-- | src/jlens.py | 9 |
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] |