Add run name as wandb config

talmolab · Dec 13, 2024 · 0e6b97c · 0e6b97c
1 parent d675b0d
commit 0e6b97c
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/sleap_nn/training/model_trainer.py b/sleap_nn/training/model_trainer.py
@@ -681,17 +681,18 @@ def train(
         )
 
         try:
+            if self.config.trainer_config.use_wandb:
+                wandb_logger.experiment.config.update({"run_name": wandb_config.name})
+                wandb_logger.experiment.config.update(dict(self.config))
+                wandb_logger.experiment.config.update({"model_params": total_params})
+
             self.trainer.fit(
                 self.model,
                 self.train_data_loader,
                 self.val_data_loader,
                 ckpt_path=self.config.trainer_config.resume_ckpt_path,
             )
 
-            if self.config.trainer_config.use_wandb:
-                wandb_logger.experiment.config.update(dict(self.config))
-                wandb_logger.experiment.config.update({"model_params": total_params})
-
         except KeyboardInterrupt:
             print("Stopping training...")