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