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

chess-gpt-eval-contrastive/mamba_module.py

@@ -110,6 +110,7 @@ class MambaPlayer:
             else:
                 self.linear_probes = {}
         if update_contrastive or update_linear:
+            linear_size = self.model.config.d_model * self.max_seq_len
             for i, layer in enumerate(self.model.backbone.layers):
                 self.activations_sum[i] = {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)),
@@ -132,9 +133,9 @@ class MambaPlayer:
                 if update_linear:
                     if not linear_probe_path or not os.path.exists(linear_probe_path):
                         self.linear_probes[i] = {
-                            'q_value': nn.Linear(self.model.config.d_model, 1),
-                            'q_value_delta': nn.Linear(self.model.config.d_model, 1),
-                            'material_balance': nn.Linear(self.model.config.d_model, 1)
+                            'q_value': nn.Linear(linear_size, 1),
+                            'q_value_delta': nn.Linear(linear_size, 1),
+                            'material_balance': nn.Linear(linear_size, 1)
                         }
             if update_linear:
                 self.linear_probe_targets = {i: {bucket: {'q_value': [], 'q_value_delta': [], 'material_balance': []} for bucket in self.move_buckets} for i in self.linear_probes}
@@ -331,7 +332,7 @@ class MambaPlayer:
         for layer_idx in self.linear_probes:
             for bucket in self.move_buckets:
                 if self.activations_count[layer_idx][bucket]['current'] > 0:
-                    X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float() #/ self.activations_count[layer_idx][bucket]['current']).float()
+                    X = torch.from_numpy(self.activations_sum[layer_idx][bucket]['current']).float().flatten(1) #/ self.activations_count[layer_idx][bucket]['current']).float()
                     for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
                         y = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().unsqueeze(1)
                         if len(y) > 0: