set env var TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 for pl < 1.5.0

Author: matsumotosan

src/lightning/fabric/utilities/cloud_io.py

@@ -26,7 +26,6 @@
 from fsspec.implementations.local import AbstractFileSystem
 from lightning_utilities.core.imports import module_available
 
-from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 
 log = logging.getLogger(__name__)
@@ -49,11 +48,6 @@ def _load(
             `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`_.
 
     """
-    # default to `weights_only=True` for torch>=2.6
-    if weights_only is None and _TORCH_GREATER_EQUAL_2_6:
-        log.debug("Defaulting to `weights_only=True` for torch>=2.6.")
-        weights_only = True
-
     if not isinstance(path_or_url, (str, Path)):
         # any sort of BytesIO or similar
         return torch.load(

tests/tests_pytorch/checkpointing/test_legacy_checkpoints.py

@@ -105,9 +105,14 @@ def load_model():
 
 @pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
 @RunIf(sklearn=True)
-def test_resume_legacy_checkpoints(tmp_path, pl_version: str):
+def test_resume_legacy_checkpoints(monkeypatch, tmp_path, pl_version: str):
     PATH_LEGACY = os.path.join(LEGACY_CHECKPOINTS_PATH, pl_version)
     with patch("sys.path", [PATH_LEGACY] + sys.path):
+        if pl_version == "local":
+            pl_version = pl.__version__
+        if Version(pl_version) < Version("1.5.0"):
+            monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")
+
         path_ckpts = sorted(glob.glob(os.path.join(PATH_LEGACY, f"*{CHECKPOINT_EXTENSION}")))
         assert path_ckpts, f'No checkpoints found in folder "{PATH_LEGACY}"'
         path_ckpt = path_ckpts[-1]