Torch: test on 2.0 and latest versions + explicitly load with weights_only=True

.github/workflows/python-tests.yml

@@ -26,7 +26,8 @@ jobs:
           [
             "Repository only",
             "Everything else",
-            "torch",
+            "torch_1.11",
+            "torch_latest",
           ]
         include:
           - python-version: "3.11" # LFS not ran on 3.8
@@ -71,8 +72,18 @@ jobs:
               git config --global user.name "ci"
               ;;
 
-            fastai | torch)
-              uv pip install "huggingface_hub[${{ matrix.test_name }}] @ ."
+            fastai)
+              uv pip install "huggingface_hub[fastai] @ ."
+              ;;
+
+            torch_latest)
+              uv pip install "huggingface_hub[torch] @ ."
+              uv pip install --upgrade torch
+              ;;
+
+            torch_1.11)
+              uv pip install "huggingface_hub[torch] @ ."
+              uv pip install torch~=1.11
               ;;
 
             tensorflow)
@@ -121,7 +132,7 @@ jobs:
               eval "$PYTEST ../tests/test_serialization.py"
             ;;
 
-            torch)
+            torch_1.11 | torch_latest)
             eval "$PYTEST ../tests/test_hub_mixin*"
             eval "$PYTEST ../tests/test_serialization.py"
             ;;

src/huggingface_hub/hub_mixin.py

@@ -820,7 +820,7 @@ def _from_pretrained(
 
     @classmethod
     def _load_as_pickle(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        state_dict = torch.load(model_file, map_location=torch.device(map_location))
+        state_dict = torch.load(model_file, map_location=torch.device(map_location), weights_only=True)
         model.load_state_dict(state_dict, strict=strict)  # type: ignore
         model.eval()  # type: ignore
         return model

tests/test_hub_mixin_pytorch.py

@@ -13,6 +13,7 @@
 from huggingface_hub.constants import PYTORCH_WEIGHTS_NAME
 from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
 from huggingface_hub.hub_mixin import ModelHubMixin, PyTorchModelHubMixin
+from huggingface_hub.serialization._torch import storage_ptr
 from huggingface_hub.utils import SoftTemporaryDirectory, is_torch_available
 
 from .testing_constants import ENDPOINT_STAGING, TOKEN, USER
@@ -415,10 +416,10 @@ def forward(self, x):
 
         # Linear layers should share weights and biases in memory
         state_dict = reloaded.state_dict()
-        a_weight_ptr = state_dict["a.weight"].untyped_storage().data_ptr()
-        b_weight_ptr = state_dict["b.weight"].untyped_storage().data_ptr()
-        a_bias_ptr = state_dict["a.bias"].untyped_storage().data_ptr()
-        b_bias_ptr = state_dict["b.bias"].untyped_storage().data_ptr()
+        a_weight_ptr = storage_ptr(state_dict["a.weight"])
+        b_weight_ptr = storage_ptr(state_dict["b.weight"])
+        a_bias_ptr = storage_ptr(state_dict["a.bias"])
+        b_bias_ptr = storage_ptr(state_dict["b.bias"])
         assert a_weight_ptr == b_weight_ptr
         assert a_bias_ptr == b_bias_ptr