Address mypy warning

microsoft · Dec 13, 2024 · 07c270a · 07c270a
1 parent b64f222
commit 07c270a
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/onnxscript/rewriter/onnxruntime/xformers/_test_models.py b/onnxscript/rewriter/onnxruntime/xformers/_test_models.py
@@ -70,8 +70,9 @@ def get_onnx_model(self):
             model, inputs, input_names=input_names, dynamo=True, fallback=True
         )
         # ORT Transformer optimizations are applied after basic optimization.
-        onnxscript.optimizer.optimize(exported.model)  # type: ignore[union-attr]
-        return exported.model
+        exported_model = exported.model  # type: ignore[union-attr]
+        onnxscript.optimizer.optimize(exported_model)
+        return exported_model
 
     def get_inputs(self):
         if not hasattr(self, "_inputs"):