import

Author: Borda

src/lightning/pytorch/core/module.py

@@ -61,7 +61,7 @@
 from lightning.pytorch.trainer.connectors.logger_connector.result import _get_default_dtype
 from lightning.pytorch.utilities import GradClipAlgorithmType
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from lightning.pytorch.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_0_9_1
+from lightning.pytorch.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_0_9_1, _TORCH_GREATER_EQUAL_2_6
 from lightning.pytorch.utilities.model_helpers import _restricted_classmethod
 from lightning.pytorch.utilities.rank_zero import WarningCache, rank_zero_warn
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
@@ -73,14 +73,16 @@
     OptimizerLRScheduler,
 )
 
-if TYPE_CHECKING:
-    from torch.distributed.device_mesh import DeviceMesh
-
 _ONNX_AVAILABLE = RequirementCache("onnx")
 _ONNXSCRIPT_AVAILABLE = RequirementCache("onnxscript")
 
-if TYPE_CHECKING and _ONNXSCRIPT_AVAILABLE:
-    from torch.onnx import ONNXProgram
+if TYPE_CHECKING:
+    from torch.distributed.device_mesh import DeviceMesh
+
+    if _TORCH_GREATER_EQUAL_2_6:
+        from torch.onnx import ONNXProgram
+    else:
+        from torch.onnx._internal.exporter import ONNXProgram
 
 warning_cache = WarningCache()
 log = logging.getLogger(__name__)

src/lightning/pytorch/utilities/imports.py

@@ -14,13 +14,15 @@
 """General utilities."""
 
 import functools
+import operator
 import sys
 
-from lightning_utilities.core.imports import RequirementCache, package_available
+from lightning_utilities.core.imports import RequirementCache, package_available, compare_version
 
 from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 
 _PYTHON_GREATER_EQUAL_3_11_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 11)
+_TORCH_GREATER_EQUAL_2_6 = compare_version("torch", operator.ge, "2.6.0")
 _TORCHMETRICS_GREATER_EQUAL_0_8_0 = RequirementCache("torchmetrics>=0.8.0")
 _TORCHMETRICS_GREATER_EQUAL_0_9_1 = RequirementCache("torchmetrics>=0.9.1")
 _TORCHMETRICS_GREATER_EQUAL_0_11 = RequirementCache("torchmetrics>=0.11.0")  # using new API with task

tests/tests_pytorch/models/test_onnx.py

@@ -25,12 +25,18 @@
 from lightning_utilities import compare_version
 
 import tests_pytorch.helpers.pipelines as tpipes
+from lightning.pytorch.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.pytorch import Trainer
 from lightning.pytorch.core.module import _ONNXSCRIPT_AVAILABLE
 from lightning.pytorch.demos.boring_classes import BoringModel
 from tests_pytorch.helpers.runif import RunIf
 from tests_pytorch.utilities.test_model_summary import UnorderedModel
 
+if _TORCH_GREATER_EQUAL_2_6:
+    from torch.onnx import ONNXProgram
+else:
+    from torch.onnx._internal.exporter import ONNXProgram
+
 
 @RunIf(onnx=True)
 def test_model_saves_with_input_sample(tmp_path):
@@ -207,14 +213,10 @@ def test_model_return_type():
     model.eval()
 
     onnx_pg = model.to_onnx(dynamo=True)
-
-    onnx_cls = torch.onnx.ONNXProgram if torch.__version__ >= "2.6.0" else torch.onnx._internal.exporter.ONNXProgram
-
-    assert isinstance(onnx_pg, onnx_cls)
+    assert isinstance(onnx_pg, ONNXProgram)
 
     model_ret = model(model.example_input_array)
     inf_ret = onnx_pg(model.example_input_array)
-
     assert torch.allclose(model_ret, inf_ret[0], rtol=1e-03, atol=1e-05)