From 2dd6d20d4c150f64d1a2f6ebca1e0bc3717aa9c8 Mon Sep 17 00:00:00 2001 From: "Wu, Jiantao (PG/R - Comp Sci & Elec Eng)" <jiantao.wu@surrey.ac.uk> Date: Wed, 14 Feb 2024 17:39:23 +0000 Subject: [PATCH] Add comment to module_args in Trainer class --- vitookit/utils/submitit.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vitookit/utils/submitit.py b/vitookit/utils/submitit.py index f936380..b275a94 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}") -- GitLab