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

chess-gpt-eval-contrastive/mamba_module.py

@@ -309,14 +309,15 @@ class MambaPlayer:
             for bucket in self.move_buckets:
                 if self.activations_count[layer_idx][bucket]['current'] > 0:
                     X = self.activations_sum[layer_idx][bucket]['current'] / self.activations_count[layer_idx][bucket]['current']
-                    X = X.reshape(X.shape[0], -1)  # Reshape X to have 2 dimensions
+                    X = X.reshape(X.shape[1], -1)  # Reshape X to have shape (sequence_length, flattened_features)
                     for probe_type in ['q_value', 'q_value_delta', 'material_balance']:
                         y = np.array(self.linear_probe_targets[layer_idx][bucket][probe_type])
-                        if len(y) == X.shape[0]:  # Check if the number of samples match
-                            self.linear_probes[layer_idx][probe_type].fit(X, y)
-                            print(f"Fit layer {layer_idx} type {probe_type}.")
+                        if len(y) > 0:
+                            # Repeat X to match the number of target values
+                            X_repeated = np.repeat(X, len(y), axis=0)
+                            self.linear_probes[layer_idx][probe_type].fit(X_repeated, y)
                         else:
-                            print(f"Skipping training for layer {layer_idx}, bucket {bucket}, probe type {probe_type} due to inconsistent number of samples. X shaspe {X.shape}, Y shape {y.shape}")
+                            print(f"Skipping training for layer {layer_idx}, bucket {bucket}, probe type {probe_type} due to empty target values.")
     
         # Reset linear_probe_targets after training
         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}