slurm version of main_distributed.py

This commit is contained in:
ahyanc25@reutlingen-university.de
2026-04-18 12:28:40 +02:00
parent f513bc97b4
commit 69d811ecb4
+6 -6
View File
@@ -73,15 +73,15 @@ class Trainer:
def launch(): def launch():
executor = submitit.AutoExecutor( executor = submitit.SlurmExecutor(
folder=os.path.join(args.folder, 'job_%j'), folder=os.path.join(args.folder, 'job_%j'),
slurm_max_num_timeout=20) max_num_timeout=20)
executor.update_parameters( executor.update_parameters(
slurm_partition=args.partition, partition=args.partition,
slurm_mem_per_gpu='55G', mem_per_gpu='55G',
timeout_min=args.time, time=args.time,
nodes=args.nodes, nodes=args.nodes,
tasks_per_node=args.tasks_per_node, ntasks_per_node=args.tasks_per_node,
cpus_per_task=10, cpus_per_task=10,
gpus_per_node=args.tasks_per_node) gpus_per_node=args.tasks_per_node)