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

chess-gpt-eval-contrastive/mamba_module.py

@@ -94,8 +94,8 @@ class MambaPlayer:
         self.activations_count = {}
         
         for i, layer in enumerate(self.model.backbone.layers):
-            self.activations_sum[i] = {"won": 0, "lost": 0}
-            self.activations_count[i] = {"won": 0, "lost": 0}
+            self.activations_sum[i] = {"won": 0, "lost": 0, "current": 0}
+            self.activations_count[i] = {"won": 0, "lost": 0, "current": 0}
             
             def hook(module, input, output, layer_idx=i):
                 if isinstance(output, tuple):
@@ -199,6 +199,9 @@ class MambaPlayer:
 
         self.activations_sum = {}
         self.activations_count = {}
+        for i, layer in enumerate(self.model.backbone.layers):
+            self.activations_sum[i] = {"won": 0, "lost": 0, "current": 0}
+            self.activations_count[i] = {"won": 0, "lost": 0, "current": 0}
     
     def apply_contrastive_activations(self, path):
         if os.path.exists(path):