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

chess-gpt-eval-contrastive/mamba_module.py

@@ -10,6 +10,7 @@ import chess
 from sklearn.linear_model import LinearRegression
 import torch.nn as nn
 import torch.optim as optim
+import wandb
 
 BASE_DIR = "mamba/"
 
@@ -142,6 +143,7 @@ class MambaPlayer:
                     }
                     for layer_idx in self.linear_probes
                 }
+                wandb.init(project="mamba_linear_probes", name=f"mamba_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()
@@ -327,6 +329,7 @@ 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()})
 
         # 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}