Reflect Andrei's comments

NVIDIA · Nov 18, 2024 · fdaef60 · fdaef60
1 parent da6037c
commit fdaef60
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py b/tests/slurm_command_gen_strategy/test_nemo_run_slurm_command_gen_strategy.py
@@ -59,7 +59,7 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> NeMoRunSlurmCommandGenS
         "cmd_args, expected_cmd",
         [
             (
-                {"task": "fine_tune", "recipe_name": "llama7_13b"},
+                {"docker_image_url": "nvcr.io/nvidia/nemo:24.09", "task": "fine_tune", "recipe_name": "llama7_13b"},
                 ["nemo", "llm", "fine_tune", "--factory", "llama7_13b", "-y", "trainer.num_nodes=2", "extra_args"],
             ),
         ],
@@ -83,7 +83,7 @@ def test_generate_test_command(
     @pytest.mark.parametrize(
         "cmd_args, expected_exception",
         [
-            ({"recipe_name": "llama7_13b"}, ValueError),
+            ({"docker_image_url": "nvcr.io/nvidia/nemo:24.09", "recipe_name": "llama7_13b"}, ValueError),
             ({"task": "fine_tune"}, ValueError),
         ],
     )