Enable optimization via environment variable

Signed-off-by: Ganesan Ramalingam <[email protected]>
microsoft · Sep 17, 2024 · 6b0f78c · 6b0f78c
1 parent 82dac0f
commit 6b0f78c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py
@@ -19,7 +19,7 @@
 
 import onnx
 
-from onnxscript import ir
+from onnxscript import ir, optimizer
 from onnxscript.function_libs.torch_lib import registration
 from onnxscript.ir import _external_data
 
@@ -28,6 +28,9 @@
     os.getenv("TORCH_ONNX_OFFLOAD_EXTERNAL_DATA_WITH_IR") != "0"
 )
 
+# Internal flag. Will go away.
+_TORCH_ONNX_ENABLE_OPTIMIZATION = os.getenv("TORCH_ONNX_ENABLE_OPTIMIZATION") == "1'"
+
 
 @dataclasses.dataclass(frozen=True)
 class _OnnxFunctionMeta:
@@ -50,7 +53,8 @@ class _OnnxFunctionMeta:
 def optimize(model: ir.Model) -> ir.Model:
     """Optimize the model."""
 
-    # TODO(justinchuby): Use the optimizer
+    if _TORCH_ONNX_ENABLE_OPTIMIZATION:
+        optimizer.optimize_ir(model)
     return model