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

chess-gpt-eval-contrastive/mamba_module.py

@@ -383,5 +383,6 @@ class MambaPlayer:
                 probe = self.linear_probes[layer_idx][probe_type]
                 #probe.eval()
                 prediction = probe(X).item()
-                print(f"Layer {layer_idx}, {probe_type}: {prediction} vs {target}")
+                if probe_type == 'material_balance':
+                    print(f"Layer {layer_idx}, {probe_type}: {int(prediction)} vs {int(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}