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

chess-gpt-eval-contrastive/mamba_module.py

@@ -97,7 +97,8 @@ class MambaPlayer:
             def hook(module, input, output, layer_idx=i):
                 self.activations[layer_idx]["current"] = output[0]#.detach().cpu().numpy()
             
-            self.hooks.append(layer.register_forward_hook(hook))
+            hook_function = partial(hook, layer_idx=i)
+            self.hooks.append(layer.register_forward_hook(hook_function))
 
     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()