Use model

huggingface · Mar 4, 2024 · 2849871 · 2849871
1 parent 5b48a2e
commit 2849871
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
@@ -463,7 +463,7 @@ class PyTorchModelHubMixin(ModelHubMixin):
     def _save_pretrained(self, save_directory: Path) -> None:
         """Save weights from a Pytorch model to a local directory."""
         model_to_save = self.module if hasattr(self, "module") else self  # type: ignore
-        save_model(model_to_save.state_dict(), save_directory / SAFETENSORS_SINGLE_FILE)
+        save_model(model_to_save, save_directory / SAFETENSORS_SINGLE_FILE)
 
     @classmethod
     def _from_pretrained(