From 3846a62f56e4170598e05aafab79d4678584db7a Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Mon, 4 Aug 2025 15:36:08 -0700
Subject: [PATCH 1/2] Add ability to pass model_dir to
 .ci/scripts/test_huggingface_optimum_model.py

---
 .ci/scripts/test_huggingface_optimum_model.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/.ci/scripts/test_huggingface_optimum_model.py b/.ci/scripts/test_huggingface_optimum_model.py
index 8a0b244c549..031479aebca 100644
--- a/.ci/scripts/test_huggingface_optimum_model.py
+++ b/.ci/scripts/test_huggingface_optimum_model.py
@@ -270,6 +270,7 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
     parser.add_argument("--model", type=str, required=True)
     parser.add_argument("--recipe", type=str, required=True)
     parser.add_argument("--quantize", action="store_true", help="Enable quantization")
+    parser.add_argument("--model_dir", type=str, required=False)
     args = parser.parse_args()
 
     model_to_model_id_and_test_function = {
@@ -294,11 +295,20 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
             f"Unknown model name: {args.model}. Available models: {model_to_model_id_and_test_function.keys()}"
         )
 
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        model_id, test_fn = model_to_model_id_and_test_function[args.model]
+    model_id, test_fn = model_to_model_id_and_test_function[args.model]
+    if args.model_dir is None:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            test_fn(
+                model_id=model_id,
+                model_dir=tmp_dir,
+                recipe=args.recipe,
+                quantize=args.quantize,
+            )
+    else:
         test_fn(
             model_id=model_id,
-            model_dir=tmp_dir,
+            model_dir=args.model_dir,
             recipe=args.recipe,
             quantize=args.quantize,
+            run_only=False,
         )

From 4815b11429548d6ee48dbf9ff05584e06fb96974 Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Tue, 5 Aug 2025 17:31:22 -0700
Subject: [PATCH 2/2] up

---
 .ci/scripts/test_huggingface_optimum_model.py | 22 ++++++++-----------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/.ci/scripts/test_huggingface_optimum_model.py b/.ci/scripts/test_huggingface_optimum_model.py
index 031479aebca..6a31eabb0c8 100644
--- a/.ci/scripts/test_huggingface_optimum_model.py
+++ b/.ci/scripts/test_huggingface_optimum_model.py
@@ -262,7 +262,7 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
 
     assert torch.allclose(
         eager_output.logits, et_output, atol=1e-02, rtol=1e-02
-    ), "CoreML output does not match eager"
+    ), "Model output does not match eager"
 
 
 if __name__ == "__main__":
@@ -270,7 +270,12 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
     parser.add_argument("--model", type=str, required=True)
     parser.add_argument("--recipe", type=str, required=True)
     parser.add_argument("--quantize", action="store_true", help="Enable quantization")
-    parser.add_argument("--model_dir", type=str, required=False)
+    parser.add_argument(
+        "--model_dir",
+        type=str,
+        required=False,
+        help="When provided, write the pte file to this directory. Otherwise, a temporary directory is created for the test.",
+    )
     args = parser.parse_args()
 
     model_to_model_id_and_test_function = {
@@ -296,19 +301,10 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False):
         )
 
     model_id, test_fn = model_to_model_id_and_test_function[args.model]
-    if args.model_dir is None:
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            test_fn(
-                model_id=model_id,
-                model_dir=tmp_dir,
-                recipe=args.recipe,
-                quantize=args.quantize,
-            )
-    else:
+    with tempfile.TemporaryDirectory() as tmp_dir:
         test_fn(
             model_id=model_id,
-            model_dir=args.model_dir,
+            model_dir=tmp_dir if args.model_dir is None else args.model_dir,
             recipe=args.recipe,
             quantize=args.quantize,
-            run_only=False,
         )