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

chess-gpt-eval-contrastive/mamba_module.py

@@ -203,9 +203,17 @@ class MambaPlayer:
 
     def update_activations(self, result):
         for layer_idx in self.activations_sum:
-            for bucket in self.move_buckets:
-                self.activations_sum[layer_idx][bucket][result] += self.activations_sum[layer_idx][bucket]["current"]
-                self.activations_count[layer_idx][bucket][result] += 1
+            if "result" == "reset":
+                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}
+            else:
+                for bucket in self.move_buckets:
+                    self.activations_sum[layer_idx][bucket][result] += self.activations_sum[layer_idx][bucket]["current"]
+                    self.activations_count[layer_idx][bucket][result] += 1
     
     def save_activations(self, path):
         if os.path.exists(path):