Adjust --gres=gpu:8 to number of user specified devices

sean-smith · sean-smith · commit ea013a7a012b · 2025-08-25T16:36:26.000-07:00
Signed-off-by: Sean Smith &lt;seasmith@nvidia.com&gt;
diff --git a/large_language_model_pretraining/nemo/pretrain_llama31.py b/large_language_model_pretraining/nemo/pretrain_llama31.py
@@ -75,7 +75,7 @@ def slurm_executor(
         gpus_per_node=devices,
         mem="0",
         exclusive=True,
-        gres="gpu:8",
+        gres=f"gpu:{devices}",
         packager=run.GitArchivePackager(),
         dependencies=dependencies,
     )

Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ def slurm_executor(`
`75`	`75`	`gpus_per_node=devices,`
`76`	`76`	`mem="0",`
`77`	`77`	`exclusive=True,`
`78`		`- gres="gpu:8",`
	`78`	`+ gres=f"gpu:{devices}",`
`79`	`79`	`packager=run.GitArchivePackager(),`
`80`	`80`	`dependencies=dependencies,`
`81`	`81`	`)`