Update test_onnxruntime_genai_e2e.py

microsoft · Apr 1, 2024 · 8c4d3ca · 8c4d3ca
1 parent 1c5fe02
commit 8c4d3ca
Showing 1 changed file with 5 additions and 7 deletions.
diff --git 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)