Update encode.py for single-node

showlab · Jul 30, 2024 · 0299f10 · 0299f10
1 parent 26581a2
commit 0299f10
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/data/preprocess/encode.py b/data/preprocess/encode.py
@@ -24,7 +24,7 @@ class LiveOnePlusEncodingArguments(LiveOnePlusTrainingArguments):
         vision_encode=vision_encode, 
         batch_size=256, save_bf16=True
     )
-    executor = submitit.AutoExecutor(folder=f"outputs/preprocess/")
+    executor = submitit.AutoExecutor(folder=f"outputs/preprocess/", cluster='local' if args.num_nodes == 1 else 'slurm')
     executor.update_parameters(
         tasks_per_node=args.num_gpus,
         nodes=args.num_nodes,
@@ -33,5 +33,6 @@ class LiveOnePlusEncodingArguments(LiveOnePlusTrainingArguments):
         slurm_partition=args.slurm_partition,
         mem_gb=240,
         slurm_time='24:00:00',
+        timeout_min=600,
     )
     job = executor.submit(task)