diff --git a/vitookit/utils/submitit.py b/vitookit/utils/submitit.py index f936380a99f50d7761b95f64974833ea77926c00..b275a94343f37d363d34fbcd7c3fc69c7c5995e0 100644 --- a/vitookit/utils/submitit.py +++ b/vitookit/utils/submitit.py @@ -113,6 +113,8 @@ class Trainer(object): module_args.gpu = job_env.local_rank module_args.rank = job_env.global_rank module_args.world_size = job_env.num_tasks + + module_args.comment = f"Job {job_env.job_id} on {job_env.num_tasks} GPUs" print("Setting up GPU args", module_args) print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")