Addressing comments

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

Author: i-riyad

modelopt/torch/quantization/export_onnx.py

@@ -107,12 +107,6 @@
 import onnx
 import torch
 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
@@ -297,6 +291,11 @@ def scaled_dot_product_attention(
     enable_gqa: bool = False,
 ):
     """Perform scaled dot product attention."""
+    if Version(torch.__version__) > Version("2.8.0"):
+        from torch.onnx._internal.torchscript_exporter import _type_utils
+    else:
+        from torch.onnx import _type_utils
+
     assert (not is_causal) or (is_causal and symbolic_helper._is_none(attn_mask)), (
         "is_causal and attn_mask cannot be set at the same time"
     )
@@ -400,6 +399,11 @@ def export_fp8_mha(
     """
     from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask
 
+    if Version(torch.__version__) > Version("2.8.0"):
+        from torch.onnx._internal.torchscript_exporter import _type_utils
+    else:
+        from torch.onnx import _type_utils
+
     # Pass all arguments, including x, to the custom ONNX operator
     assert (not is_causal) or (is_causal and sym_help._is_none(attn_mask)), (
         "is_causal and attn_mask cannot be set at the same time"

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

@@ -32,12 +32,6 @@
 from packaging.version import Version
 from torch import nn
 
-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
 
@@ -885,6 +879,11 @@ def forward(self, inputs):
         Returns:
             outputs: A Tensor of type output_dtype
         """
+        if Version(torch.__version__) > Version("2.8.0"):
+            from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
+        else:
+            from torch.onnx._globals import GLOBALS
+
         if DTensor is not None and isinstance(inputs, DTensor):
             # TensorQuantizer only handles regular non-DTensor inputs
             device_mesh, placements = inputs.device_mesh, inputs.placements