Reflect Andrei's comments

NVIDIA · Nov 19, 2024 · 12122d5 · 12122d5
1 parent 17fa827
commit 12122d5
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/cloudai/schema/test_template/nemo_run/slurm_command_gen_strategy.py b/src/cloudai/schema/test_template/nemo_run/slurm_command_gen_strategy.py
@@ -36,12 +36,9 @@ def _parse_slurm_args(
         return base_args
 
     def generate_test_command(self, env_vars: Dict[str, str], cmd_args: Dict[str, str], tr: TestRun) -> List[str]:
-        command = ["nemo", "llm"]
-
         tdef: NeMoRunTestDefinition = cast(NeMoRunTestDefinition, tr.test.test_definition)
-        command.append(tdef.cmd_args.task)
-        command.extend(["--factory", tdef.cmd_args.recipe_name])
-        command.append("-y")
+
+        command = ["nemo", "llm", tdef.cmd_args.task, "--factory", tdef.cmd_args.recipe_name, "-y"]
 
         if tr.nodes:
             command.append(f"trainer.num_nodes={len(tr.nodes)}")