From 0fff2420988e9f1fcbb557d416a7eb7c4f098d71 Mon Sep 17 00:00:00 2001
From: jinchen62 <jinchenye62@gmail.com>
Date: Tue, 20 Aug 2024 13:56:34 -0700
Subject: [PATCH] Fix shape inference when importing models

---
 e2eshark/run.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/e2eshark/run.py b/e2eshark/run.py
index 8742282f..c0729944 100644
--- a/e2eshark/run.py
+++ b/e2eshark/run.py
@@ -120,6 +120,7 @@ def runOnnxToTorchMLIRGeneration(
         logfilename = curphase + ".log"
         scriptcommand = (
             f"{torch_mlir_pythonpath} python -m torch_mlir.tools.import_onnx "
+            + "--opset-version=21 "
             + onnxfilename
             + " -o "
             + torchonnxfilename