Add support for torch.load with weights_only=True

Author: stes

cebra/integrations/sklearn/cebra.py

@@ -28,9 +28,9 @@
 import numpy as np
 import numpy.typing as npt
 import pkg_resources
+import sklearn
 import sklearn.utils.validation as sklearn_utils_validation
 import torch
-import sklearn
 from sklearn.base import BaseEstimator
 from sklearn.base import TransformerMixin
 from sklearn.utils.metaestimators import available_if
@@ -43,12 +43,21 @@
 import cebra.models
 import cebra.solver
 
+# NOTE(stes): From torch 2.6 onwards, we need to specify the following list
+# when loading CEBRA models to allow weights_only = True.
+CEBRA_LOAD_SAFE_GLOBALS = [
+    cebra.data.Offset, torch.torch_version.TorchVersion, np.dtype,
+    np.dtypes.Float64DType, np.dtypes.Int64DType
+]
+
+
 def check_version(estimator):
     # NOTE(stes): required as a check for the old way of specifying tags
     # https://github.com/scikit-learn/scikit-learn/pull/29677#issuecomment-2334229165
     from packaging import version
     return version.parse(sklearn.__version__) < version.parse("1.6.dev")
 
+
 def _init_loader(
     is_cont: bool,
     is_disc: bool,
@@ -1409,12 +1418,18 @@ def save(self,
     def load(cls,
              filename: str,
              backend: Literal["auto", "sklearn", "torch"] = "auto",
+             weights_only: bool = True,
              **kwargs) -> "CEBRA":
         """Load a model from disk.
 
         Args:
             filename: The path to the file in which to save the trained model.
             backend: A string identifying the used backend.
+            weights_only: Indicates whether unpickler should be restricted to loading only tensors, primitive types,
+                dictionaries and any types added via `py:func:torch.serialization.add_safe_globals`.
+                See `py:func:torch.load` with ``weights_only=True`` for more details. It it recommended to leave this
+                at the default value of ``True``. If you experience issues with loading custom models (specified outside
+                of the CEBRA package), you can try to set this to ``False`` if you trust the source of the model.
             kwargs: Optional keyword arguments passed directly to the loader.
 
         Return:
@@ -1443,14 +1458,17 @@ def load(cls,
             >>> tmp_file.unlink()
 
         """
-
         supported_backends = ["auto", "sklearn", "torch"]
         if backend not in supported_backends:
             raise NotImplementedError(
                 f"Unsupported backend: '{backend}'. Supported backends are: {', '.join(supported_backends)}"
             )
 
-        checkpoint = torch.load(filename, **kwargs)
+        if not weights_only:
+            checkpoint = torch.load(filename, weights_only=False, **kwargs)
+        else:
+            with torch.serialization.safe_globals(CEBRA_LOAD_SAFE_GLOBALS):
+                checkpoint = torch.load(filename, weights_only=True, **kwargs)
 
         if backend == "auto":
             backend = "sklearn" if isinstance(checkpoint, dict) else "torch"