diff --git a/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py b/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py
index 903ca751..3c7305ac 100644
--- a/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py
+++ b/src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py
@@ -60,6 +60,9 @@ def gen_exec_command(
         for key, value in final_env_vars.items():
             self.final_cmd_args[f"env_vars.{key}"] = value
 
+        if "training.values" in self.final_cmd_args:
+            self.final_cmd_args["training"] = self.final_cmd_args.pop("training.values")
+
         self.final_cmd_args["cluster.partition"] = self.slurm_system.default_partition
         reservation_key = "--reservation "
         if self.slurm_system.extra_srun_args and reservation_key in self.slurm_system.extra_srun_args: