From c38c2cb372be9ab46da8dfd8a5a2f81f15e70076 Mon Sep 17 00:00:00 2001 From: shubhambhokare1 Date: Thu, 21 Nov 2024 02:50:48 +0000 Subject: [PATCH] make flag dynamic --- onnxscript/_framework_apis/torch_2_5.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py index 484443f11..4fc6fda24 100644 --- a/onnxscript/_framework_apis/torch_2_5.py +++ b/onnxscript/_framework_apis/torch_2_5.py @@ -21,8 +21,6 @@ from onnxscript.function_libs.torch_lib import registration from onnxscript.ir import _external_data -ENABLE_VERSION_CONVERTER = 0 - @dataclasses.dataclass(frozen=True) class _OnnxFunctionMeta: @@ -53,7 +51,9 @@ def optimize(model: ir.Model) -> ir.Model: def convert_version(model: ir.Model, target_version: int) -> ir.Model: """Convert the model to the specified ONNX opset version.""" - if ENABLE_VERSION_CONVERTER == 1: + # Internal flag. Will go away. + enabled = os.getenv("TORCH_ONNX_ENABLE_VERSION_CONVERSION") == "1" + if enabled: version_converter.convert_version(model, target_version) return model