feat: add wandb_dir as cfg argument

alex-petrenko · Nov 18, 2024 · 8f3bb78 · 8f3bb78
1 parent 6810db5
commit 8f3bb78
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/sample_factory/cfg/cfg.py b/sample_factory/cfg/cfg.py
@@ -744,6 +744,13 @@ def add_wandb_args(p: ArgumentParser):
         nargs="*",
         help="Tags can help with finding experiments in WandB web console",
     )
+    p.add_argument(
+        "--wandb_dir",
+        default="wandb",
+        type=str,
+        nargs="*",
+        help="Logging Directory for WandB",
+    )
 
 
 def add_pbt_args(p: ArgumentParser):

diff --git a/sample_factory/utils/wandb_utils.py b/sample_factory/utils/wandb_utils.py
@@ -44,6 +44,7 @@ def init_wandb_func():
             tags=cfg.wandb_tags,
             resume="allow",
             settings=wandb.Settings(start_method="fork"),
+            dir=cfg.wandb_dir,
         )
 
     log.debug("Initializing WandB...")