Update model_utils.py

model_utils.py

@@ -18,27 +18,29 @@ def run_activation_patching(prompt):
     # Load HookedTransformer
     model = HookedTransformer.from_pretrained("gpt2-small")
 
-    # Tokenize
+    # Tokenize the text
     tokens = model.to_tokens(prompt)
 
-    # Store activations here
+    # Dict to store activations
     activations = {}
 
     # Hook function
-    def hook_fn_hook(value, hook):
+    def hook_fn(value, hook):
         activations[hook.name] = value.detach().cpu().numpy()
 
-    # Register hooks using new `.hooks` API
-    hooks = []
+    # Register hooks for all MLP post-activations
+    hook_handles = []
     for i in range(model.cfg.n_layers):
         hook_name = f"blocks.{i}.mlp.hook_post"
-        hooks.append(model.hooks.add_hook(hook_name, hook_fn_hook))
+        handle = model.add_hook(hook_name, hook_fn)
+        hook_handles.append(handle)
 
-    # Run model forward
+    # Run the model forward
     _ = model(tokens)
 
-    # Remove hooks
-    for h in hooks:
+    # Remove all hooks
+    for h in hook_handles:
         h.remove()
 
     return activations
+