Making conditional import for torch 2.9

Signed-off-by: Riyad Islam <[email protected]>

Author: i-riyad

modelopt/torch/quantization/export_onnx.py

@@ -106,7 +106,14 @@
 
 import onnx
 import torch
-from torch.onnx import _type_utils, symbolic_helper
+from packaging.version import Version
+
+if Version(torch.__version__) >= Version("2.9.0"):
+    from torch.onnx._internal.torchscript_exporter import _type_utils
+else:
+    from torch.onnx import _type_utils
+
+from torch.onnx import symbolic_helper
 from torch.onnx import symbolic_helper as sym_help
 from torch.onnx._internal import jit_utils
 from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask

modelopt/torch/quantization/nn/modules/tensor_quantizer.py

@@ -31,7 +31,12 @@
 import torch.nn.functional as F
 from packaging.version import Version
 from torch import nn
-from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
+
+if Version(torch.__version__) >= Version("2.9.0"):
+    from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
+else:
+    from torch.onnx._globals import GLOBALS
+
 
 from modelopt.torch.utils import standardize_constructor_args
 from modelopt.torch.utils.distributed import DistributedProcessGroup