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

chess-gpt-eval-contrastive/mamba_module.py

@@ -337,7 +337,6 @@ class MambaPlayer:
                         y = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().unsqueeze(1)
                         if len(y) > 0:
                             y_pred = self.linear_probes[layer_idx][probe_type](X)
-                            print(f'{y_pred}')
                             loss = criterion(y_pred, y)
                             for param_group in self.linear_optimizers[layer_idx][probe_type].param_groups:
                                 param_group['lr'] = lr
@@ -364,11 +363,11 @@ 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()
+            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().unsqueeze(1)
+                target = torch.tensor(self.linear_probe_targets[layer_idx][bucket][probe_type]).float().item()
                 probe = self.linear_probes[layer_idx][probe_type]
                 probe.eval()
-                prediction = probe(X)#.item()
+                prediction = probe(X).item()
                 print(f"Layer {layer_idx}, {probe_type}: {prediction} vs {target}")
         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}