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

chess-gpt-eval-contrastive/mamba_module.py

@@ -170,11 +170,9 @@ class MambaPlayer:
         return {"model": self.model_name}
 
     def update_activations(self, result):
-        for layer_idx in self.activations_sum:
-            self.activations_sum[layer_idx][result] += self.activations_sum[layer_idx]["current"]
-            self.activations_count[layer_idx][result] += self.activations_count[layer_idx]["current"]
-            self.activations_sum[layer_idx]["current"] = 0
-            self.activations_count[layer_idx]["current"] = 0
+        for layer_idx in self.activations:
+            self.activations_sum[layer_idx][result] += self.activations[layer_idx]["current"]
+            self.activations_count[layer_idx][result] += 1
 
     def save_activations(self, path):
         if os.path.exists(path):