Support W&B resume (NVlabs#145)

gheinrich · Jul 30, 2024 · d21e90a · d21e90a
1 parent 4fced82
commit d21e90a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/setups/train.sh b/scripts/setups/train.sh
@@ -17,8 +17,10 @@ OUTPUT_DIR=${OUTPUT_DIR:-"runs/$RUN_NAME"}
 echo "OUTPUT_DIR = $OUTPUT_DIR"
 
 export WANDB_PROJECT="vila"
-export WANDB_NAME=$RUN_NAME
 export WANDB_DIR=$OUTPUT_DIR
+export WANDB_RUN_ID=$RUN_NAME
+export WANDB_NAME=$RUN_NAME
+export WANDB_RESUME="allow"
 
 NNODES=${SLURM_JOB_NUM_NODES:-1}
 echo "NNODES = $NNODES"