update wandb resume

torchmd · Jun 19, 2023 · 4f82889 · 4f82889
1 parent b02ef20
commit 4f82889
Showing 1 changed file with 8 additions and 13 deletions.
diff --git a/torchmdnet/scripts/train.py b/torchmdnet/scripts/train.py
@@ -98,8 +98,7 @@ def get_args():
     parser.add_argument('--wandb-use', default=False, type=bool, help='Defines if wandb is used or not')
     parser.add_argument('--wandb-name', default='training', type=str, help='Give a name to your wandb run')
     parser.add_argument('--wandb-project', default='training_', type=str, help='Define what wandb Project to log to')
-    parser.add_argument('--wandb-resume', default=False, type=bool, help='Defines if wandb run should be resumed')
-    parser.add_argument('--wandb-id', default=None, type=str, help='Define what wandb run to resume, you can find the id from the wandb table')
+    parser.add_argument('--wandb-resume-from-id', default=None, type=str, help='Resume a wandb run from a given run id, retrieve the run id from the wandb dashboard')
     parser.add_argument('--tensorboard-use', default=False, type=bool, help='Defines if tensor board is used or not')
 
     # fmt: on
@@ -148,17 +147,13 @@ def main():
     csv_logger = CSVLogger(args.log_dir, name="", version="")
     _logger = [csv_logger]
     if args.wandb_use:
-        if args.wandb_resume and args.wandb_id is not None:
-            wandb_logger = WandbLogger(
-                project=args.wandb_project,
-                save_dir=args.log_dir,
-                resume="must",
-                id=args.wandb_id,
-            )
-        else:
-            wandb_logger = WandbLogger(
-                project=args.wandb_project, name=args.wandb_name, save_dir=args.log_dir
-            )
+        wandb_logger = WandbLogger(
+            project=args.wandb_project,
+            name=args.wandb_name,
+            save_dir=args.log_dir,
+            resume="must" if args.wandb_resume_from_id is not None else None,
+            id=args.wandb_resume_from_id,
+        )
         _logger.append(wandb_logger)
 
     if args.tensorboard_use: