Update chess-gpt-eval-contrastive/mamba_module.py

chess-gpt-eval-contrastive/mamba_module.py

@@ -93,7 +93,7 @@ class MambaPlayer:
         self.hooks = []
         
         for i, layer in enumerate(self.model.backbone.layers):
-            self.activations[i] = {"won": [], "lost": [], "current": []}
+            self.activations[i] = {"won": [], "lost": []}
             
             def hook(module, input, output, layer_idx=i):
                 if isinstance(output, tuple):
@@ -102,8 +102,7 @@ class MambaPlayer:
                     tensor_output = output
                 self.activations[layer_idx]["current"] = tensor_output.detach().cpu().numpy()
             
-            hook_function = partial(hook, layer_idx=i)
-            self.hooks.append(layer.register_forward_hook(hook_function))
+            self.hooks.append(layer.register_forward_hook(hook))
 
     def get_mamba_response(self, game_state: str, temperature: float, max_new_tokens: int, top_k: int):
         game_state = game_state.split("\n\n")[-1].strip()