Use 2048 args for input len 2048 llama test

Signed-off-by: aviator19941 <[email protected]>
nod-ai · Dec 5, 2024 · 6efc21c · 6efc21c
1 parent 2304a56
commit 6efc21c
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,
         )