handle device settings for ET export (#732)

* handle device settings for ET export * tab/spc
pytorch · Jul 17, 2024 · 706a5a1 · 706a5a1
1 parent b217d4b
commit 706a5a1
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/export.py b/export.py
@@ -48,7 +48,12 @@ def main(args):
     output_pte_path = args.output_pte_path
     output_dso_path = args.output_dso_path
 
-    if output_dso_path and "mps" in builder_args.device:
+    if output_pte_path and builder_args.device != "cpu":
+        print(
+            f"Warning! ExecuTorch export target is controlled by export recipe, not device setting. Ignoring device={builder_args.device} setting."
+        )
+        builder_args.device = "cpu"
+    elif "mps" in builder_args.device:
         print("Warning! Device MPS not supported for export. Exporting for device CPU.")
         builder_args.device = "cpu"