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

chess-gpt-eval-contrastive/mamba_module.py

@@ -106,13 +106,6 @@ class MambaPlayer:
                 self.linear_probes = torch.load(linear_probe_path)
             else:
                 self.linear_probes = {}
-            self.linear_optimizers = {
-                layer_idx: {
-                    probe_type: optim.Adam(self.linear_probes[layer_idx][probe_type].parameters(), lr=lr)
-                    for probe_type in ['q_value', 'q_value_delta', 'material_balance']
-                }
-                for layer_idx in self.linear_probes
-            }
         if update_contrastive or update_linear:
             for i, layer in enumerate(self.model.backbone.layers):
                 self.activations_sum[i] = {bucket: {"won": np.zeros((1, self.max_seq_len, self.model.config.d_model)),
@@ -142,6 +135,13 @@ class MambaPlayer:
                         }
             if update_linear:
                 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)
+                        for probe_type in ['q_value', 'q_value_delta', 'material_balance']
+                    }
+                    for layer_idx in self.linear_probes
+                }
 
     def get_mamba_response(self, game_state: str, temperature: float, max_new_tokens: int, top_k: int):
         game_state = game_state.split("\n\n")[-1].strip()