combine two wandb init functions

mllam · May 25, 2024 · 57a396c · 57a396c
1 parent 72272bc
commit 57a396c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 10 deletions.
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
@@ -271,10 +271,7 @@ def init_wandb_metrics(wandb_logger):
     """
     Set up wandb metrics to track
     """
-    experiment = wandb_logger.experiment
-    experiment.define_metric("val_mean_loss", summary="min")
-    for step in constants.VAL_STEP_LOG_ERRORS:
-        experiment.define_metric(f"val_loss_unroll{step}", summary="min")
+
 
 
 @rank_zero_only
@@ -303,7 +300,6 @@ def init_wandb(args):
             project=constants.WANDB_PROJECT,
             name=run_name,
             config=args,
-            log_model=True,
         )
         wandb.save("neural_lam/constants.py")
     else:
@@ -317,7 +313,10 @@ def init_wandb(args):
             project=constants.WANDB_PROJECT,
             id=args.resume_run,
             config=args,
-            log_model=True,
         )
+    experiment = logger.experiment
+    experiment.define_metric("val_mean_loss", summary="min")
+    for step in constants.VAL_STEP_LOG_ERRORS:
+        experiment.define_metric(f"val_loss_unroll{step}", summary="min")
 
     return logger
diff --git a/train_model.py b/train_model.py
@@ -281,10 +281,6 @@ def main():
         precision=args.precision,
     )
 
-    # Only init once, on rank 0 only
-    if trainer.global_rank == 0:
-        utils.init_wandb_metrics(logger)  # Do after wandb.init
-
     if args.eval:
         if args.eval == "val":
             eval_loader = val_loader