defaults for weights_only in torch.hub.load_state_dict_from_url

Author: matsumotosan

src/lightning/fabric/utilities/cloud_io.py

@@ -26,6 +26,7 @@
 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__)
@@ -56,9 +57,14 @@ def _load(
             weights_only=weights_only,
         )
     if str(path_or_url).startswith("http"):
-        if weights_only is None and _TORCH_GREATER_EQUAL_2_6:
-            weights_only = True
-            log.debug(f"Default to `weights_only=True` for remote checkpoint: {path_or_url}")
+        if weights_only is None:
+            if _TORCH_GREATER_EQUAL_2_6:
+                weights_only = True
+                log.debug(f"Default to `weights_only=True` for remote checkpoint for torch>=2.6: {path_or_url}")
+            else:
+                weights_only = False
+                log.debug(f"Default to `weights_only=False` for remote checkpoint for torch<2.6: {path_or_url}")
+
         return torch.hub.load_state_dict_from_url(
             str(path_or_url),
             map_location=map_location,  # type: ignore[arg-type]