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

chess-gpt-eval-contrastive/mamba_module.py

@@ -215,9 +215,12 @@ class MambaPlayer:
             pickle.dump((activations_sum, activations_count), f)
         
         for layer_idx in self.activations_sum:
-            for bucket in self.move_buckets:
-                self.activations_sum[layer_idx][bucket]["current"].fill(0)
-                self.activations_count[layer_idx][bucket]["current"] = 0
+            self.activations_sum[layer_idx] = {bucket: {"won": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
+                                                "lost": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
+                                                "current": np.zeros((1, self.max_seq_len, self.model.config.d_model))}
+                                       for bucket in self.move_buckets}
+            self.activations_count[layer_idx] = {bucket: {"won": 0, "lost": 0, "current": 0}
+                                         for bucket in self.move_buckets}
     
     def apply_contrastive_activations(self, path):
         if os.path.exists(path):