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

chess-gpt-eval-contrastive/mamba_module.py

@@ -97,7 +97,8 @@ class MambaPlayer:
         self.move_num = 0
         self.hooks = []
         self.max_seq_len = 1536
-        self.move_buckets = [10, 20, 30, 40, float('inf')]
+        #self.move_buckets = [10, 20, 30, 40, float('inf')]
+        self.move_buckets = [float('inf')]
         
         if update_contrastive or update_linear:
             self.activations_sum = {}
@@ -329,7 +330,8 @@ class MambaPlayer:
                             self.linear_optimizers[layer_idx][probe_type].zero_grad()
                             loss.backward()
                             self.linear_optimizers[layer_idx][probe_type].step()
-                            wandb.log({f"{probe_type}/layer_{layer_idx}_{bucket}_loss": loss.item()})
+                            #wandb.log({f"{probe_type}/layer_{layer_idx}_{bucket}_loss": loss.item()})
+                            wandb.log({f"{probe_type}/layer_{layer_idx}_loss": loss.item()})
 
         # 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}