Update ego4d_narration_refinement.py

showlab · Jul 30, 2024 · b6b5cf9 · b6b5cf9
1 parent 0299f10
commit b6b5cf9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/data/preprocess/ego4d_narration_refinement.py b/data/preprocess/ego4d_narration_refinement.py
@@ -61,7 +61,7 @@ def distributed_refine_narration(args: LiveOnePlusEncodingArguments):
 
 if __name__ == "__main__":
     args, = HfArgumentParser(LiveOnePlusEncodingArguments).parse_args_into_dataclasses()
-    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,
@@ -70,5 +70,6 @@ def distributed_refine_narration(args: LiveOnePlusEncodingArguments):
         cpus_per_task=10,
         mem_gb=240,
         slurm_time='24:00:00',
+        timeout_min=600,
     )
     job = executor.submit(distributed_refine_narration, args)