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

chess-gpt-eval-contrastive/mamba_module.py

@@ -337,8 +337,7 @@ 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'{X.shape} vs {y_pred.shape}')
-                            exit()
+                            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