Add wandb to main process

wandb · Aug 13, 2024 · 054ce76 · 054ce76
1 parent a2f9bfe
commit 054ce76
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/jobs/slurm/submitit/main.py b/jobs/slurm/submitit/main.py
@@ -3,6 +3,8 @@
 import random
 import submitit
 import time
+import wandb
+from wandb.sdk import launch
 
 def slow_multiplication(x, y):
     time.sleep(x*y)
@@ -11,6 +13,10 @@ def slow_multiplication(x, y):
 async def main():
     print(f"Running in {os.getenv("CONDA_DEFAULT_ENV")}")
     executor = submitit.AutoExecutor(folder="logs")
+    wandb.init(project="submitit-test", config={
+        "submitit": {"timeout_min": None, "partition": None}
+    })
+    launch.manage_wandb_config(include="submitit")
     # executor.update_parameters(timeout_min=1)
 
     # await a single result