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

chess-gpt-eval-contrastive/mamba_module.py

@@ -126,7 +126,7 @@ class MambaPlayer:
                         tensor_output = output
                     seq_len = tensor_output.shape[1]
                     bucket = next(b for b in self.move_buckets if self.move_num <= b)
-                    self.activations_sum[layer_idx][bucket]["current"][:, :8, :] += tensor_output.detach().cpu().numpy()[:self.seq_len][-8:]
+                    self.activations_sum[layer_idx][bucket]["current"][:, :8, :] += tensor_output.detach().cpu().numpy()[:, :self.seq_len, :][:, -8:, :]
                     self.activations_count[layer_idx][bucket]["current"] += 1
                 
                 self.hooks.append(layer.register_forward_hook(hook))
@@ -377,7 +377,7 @@ class MambaPlayer:
         self.move_num = board.fullmove_number
         bucket = next(b for b in self.move_buckets if self.move_num <= b)
         for layer_idx in self.linear_probes:
-            X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)[:self.seq_len][-8:]
+            X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1)
             for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
                 target = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().item()
                 probe = self.linear_probes[layer_idx][probe_type]