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

chess-gpt-eval-contrastive/mamba_module.py

@@ -137,7 +137,7 @@ class MambaPlayer:
                 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}
                 self.linear_optimizers = {
                     layer_idx: {
-                        probe_type: optim.Adam(self.linear_probes[layer_idx][probe_type].parameters(), lr=lr)
+                        probe_type: optim.Adam(self.linear_probes[layer_idx][probe_type].parameters(), lr=0.01)
                         for probe_type in ['q_value', 'q_value_delta', 'material_balance']
                     }
                     for layer_idx in self.linear_probes
@@ -312,7 +312,7 @@ class MambaPlayer:
             self.linear_probe_targets[layer_idx][bucket]['q_value_delta'].append(q_value_delta)
             self.linear_probe_targets[layer_idx][bucket]['material_balance'].append(material_bal)
 
-    def train_linear_probes(self, lr=0.01):
+    def train_linear_probes(self):
         criterion = nn.MSELoss()
 
         for layer_idx in self.linear_probes: