new function returns logger directly

mllam · May 1, 2024 · 96eea78 · 96eea78
1 parent 39f5976
commit 96eea78
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/train_model.py b/train_model.py
@@ -9,7 +9,7 @@
 from lightning_fabric.utilities import seed
 
 # First-party
-from neural_lam import constants, utils
+from neural_lam import utils
 from neural_lam.models.graph_lam import GraphLAM
 from neural_lam.models.hi_lam import HiLAM
 from neural_lam.models.hi_lam_parallel import HiLAMParallel
@@ -263,9 +263,9 @@ def main():
         mode="min",
         save_last=True,
     )
-    logger = pl.loggers.WandbLogger(
-        project=constants.WANDB_PROJECT, name=run_name, config=args
-    )
+
+    logger = utils.init_wandb(args)
+
     trainer = pl.Trainer(
         max_epochs=args.epochs,
         deterministic=True,