From 8c4d3caefa0f445395aa1f04f1a5aee8d1f9e9a1 Mon Sep 17 00:00:00 2001 From: "Nat Kershaw (MSFT)" Date: Mon, 1 Apr 2024 13:53:01 -0700 Subject: [PATCH] Update test_onnxruntime_genai_e2e.py --- test/python/test_onnxruntime_genai_e2e.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/test/python/test_onnxruntime_genai_e2e.py b/test/python/test_onnxruntime_genai_e2e.py index 8238c90fb..f76354261 100644 --- a/test/python/test_onnxruntime_genai_e2e.py +++ b/test/python/test_onnxruntime_genai_e2e.py @@ -51,11 +51,9 @@ def run_model(model_path: str | bytes | os.PathLike): if __name__ == "__main__": for model_name in ["microsoft/phi-2"]: - with tempfile.TemporaryDirectory() as temp_dir: - device = "cuda" if og.is_cuda_available() else "cpu" - download_model(temp_dir, device, model_name, "int4") - run_model(temp_dir) - download_model(temp_dir, "cpu", model_name, "fp32") - run_model(temp_dir) - + for precision in ["int4", "fp32"]: + with tempfile.TemporaryDirectory() as temp_dir: + device = "cuda" if og.is_cuda_available() else "cpu" + download_model(temp_dir, device, model_name, precision) + run_model(temp_dir)