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,7 +32,7 @@
 from packaging.version import Version
 from torch import nn
 
-if Version(torch.__version__) >= Version("2.9.0"):
+if Version(torch.__version__) > Version("2.8.0"):
     from torch.onnx._internal.torchscript_exporter._globals import GLOBALS
 else:
     from torch.onnx._globals import GLOBALS