From aa5d2e2f0b9e479eda1c4aa7de28d69a623c220e Mon Sep 17 00:00:00 2001
From: Github Executorch <github_executorch@arm.com>
Date: Mon, 13 Jan 2025 12:01:50 -0800
Subject: [PATCH] Fix broken QNN models on benchmarks

---
 .ci/scripts/test_model.sh         |  7 +++++++
 examples/qualcomm/scripts/edsr.py | 15 +++++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/.ci/scripts/test_model.sh b/.ci/scripts/test_model.sh
index 6e8749fe26d..b4fbc4486a2 100755
--- a/.ci/scripts/test_model.sh
+++ b/.ci/scripts/test_model.sh
@@ -169,6 +169,13 @@ test_model_with_qnn() {
     EXPORT_SCRIPT=inception_v3
   elif [[ "${MODEL_NAME}" == "vit" ]]; then
     EXPORT_SCRIPT=torchvision_vit
+  elif [[ "${MODEL_NAME}" == "edsr" ]]; then
+    EXPORT_SCRIPT=edsr
+    # Additional deps for edsr
+    pip install piq
+  else
+    echo "Unsupported model $MODEL_NAME"
+    exit 1
   fi
 
   # Use SM8450 for S22, SM8550 for S23, and SM8560 for S24
diff --git a/examples/qualcomm/scripts/edsr.py b/examples/qualcomm/scripts/edsr.py
index 517a9302e75..fa7efc0c459 100755
--- a/examples/qualcomm/scripts/edsr.py
+++ b/examples/qualcomm/scripts/edsr.py
@@ -102,14 +102,17 @@ def main(args):
             "Please specify a device serial by -s/--device argument."
         )
 
-    dataset = get_dataset(
-        args.hr_ref_dir, args.lr_dir, args.default_dataset, args.artifact
-    )
-
-    inputs, targets, input_list = dataset.lr, dataset.hr, dataset.get_input_list()
-    pte_filename = "edsr_qnn_q8"
     instance = EdsrModel()
+    if args.compile_only:
+        inputs = instance.get_example_inputs()
+    else:
+        dataset = get_dataset(
+            args.hr_ref_dir, args.lr_dir, args.default_dataset, args.artifact
+        )
 
+        inputs, targets, input_list = dataset.lr, dataset.hr, dataset.get_input_list()
+
+    pte_filename = "edsr_qnn_q8"
     build_executorch_binary(
         instance.get_eager_model().eval(),
         (inputs[0],),