Use correct args for input len 2048 llama test (#649)

Uses args for input len 2048 llama test instead of args for input len 128 llama test. --------- Signed-off-by: aviator19941 <[email protected]>
nod-ai · Dec 5, 2024 · 1763a82 · 1763a82
1 parent 8ea135a
commit 1763a82
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/tests/models/llama/benchmark_amdgpu_test.py b/sharktank/tests/models/llama/benchmark_amdgpu_test.py
@@ -270,7 +270,7 @@ def testBenchmark8B_f16_Non_Decomposed_Prefill_Input_Len_2048(self):
             hip_device_id=self.iree_device,
             vmfb_name=output_vmfb,
             irpa_path=self.irpa_path,
-            args=self.iree_run_prefill_nondecomposed_args_fp16,
+            args=self.iree_run_prefill_nondecomposed_args_fp16_2048,
             cwd=self.repo_root,
         )