refactor: update torch version checks to use greater than or equal comparisons

Author: bhimrazy

src/lightning/pytorch/utilities/imports.py

@@ -28,7 +28,7 @@
 _TORCHMETRICS_GREATER_EQUAL_0_11 = RequirementCache("torchmetrics>=0.11.0")  # using new API with task
 _TORCHMETRICS_GREATER_EQUAL_1_0_0 = RequirementCache("torchmetrics>=1.0.0")
 _TORCH_EQUAL_2_8 = RequirementCache("torch>=2.8.0,<2.9.0")
-_TORCH_EQUAL_2_9 = RequirementCache("torch>=2.9.0,<2.10.0")
+_TORCH_GREATER_EQUAL_2_8 = compare_version("torch", operator.ge, "2.8.0")
 
 _OMEGACONF_AVAILABLE = package_available("omegaconf")
 _TORCHVISION_AVAILABLE = RequirementCache("torchvision")

tests/tests_pytorch/helpers/runif.py

@@ -14,7 +14,7 @@
 import pytest
 
 from lightning.fabric.utilities.imports import _IS_WINDOWS
-from lightning.pytorch.utilities.imports import _TORCH_EQUAL_2_8, _TORCH_EQUAL_2_9
+from lightning.pytorch.utilities.imports import _TORCH_GREATER_EQUAL_2_8
 from lightning.pytorch.utilities.testing import _runif_reasons
 
 
@@ -27,6 +27,6 @@ def RunIf(**kwargs):
 _xfail_gloo_windows = pytest.mark.xfail(
     RuntimeError,
     strict=True,
-    condition=(_IS_WINDOWS and (_TORCH_EQUAL_2_8 or _TORCH_EQUAL_2_9)),
+    condition=(_IS_WINDOWS and _TORCH_GREATER_EQUAL_2_8),
     reason="makeDeviceForHostname(): unsupported gloo device",
 )

tests/tests_pytorch/trainer/test_trainer.py

@@ -55,7 +55,7 @@
 from lightning.pytorch.strategies.launchers import _MultiProcessingLauncher, _SubprocessScriptLauncher
 from lightning.pytorch.trainer.states import RunningStage, TrainerFn
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE, _TORCH_EQUAL_2_8
+from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE, _TORCH_GREATER_EQUAL_2_8
 from tests_pytorch.conftest import mock_cuda_count, mock_mps_count
 from tests_pytorch.helpers.datamodules import ClassifDataModule
 from tests_pytorch.helpers.runif import RunIf
@@ -1730,7 +1730,12 @@ def test_exception_when_lightning_module_is_not_set_on_trainer(fn):
 
 @RunIf(min_cuda_gpus=1)
 # FixMe: the memory raises to 1024 from expected 512
-@pytest.mark.xfail(AssertionError, strict=True, condition=_TORCH_EQUAL_2_8, reason="temporarily disabled for torch 2.8")
+@pytest.mark.xfail(
+    AssertionError,
+    strict=True,
+    condition=_TORCH_GREATER_EQUAL_2_8,
+    reason="temporarily disabled for torch >= 2.8",
+)
 def test_multiple_trainer_constant_memory_allocated(tmp_path):
     """This tests ensures calling the trainer several times reset the memory back to 0."""
 
@@ -1750,7 +1755,6 @@ def on_train_epoch_start(self, trainer, *_):
     def current_memory():
         # before measuring the memory force release any leftover allocations, including CUDA tensors
         gc.collect()
-        torch.cuda.empty_cache()
         return torch.cuda.memory_allocated(0)
 
     model = TestModel()