diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 702569b8c..03a624fef 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -23,8 +23,6 @@ from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack -import packaging -import safetensors from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError @@ -221,26 +219,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): @classmethod def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: - # Create base kwargs - kwargs = {"strict": strict} - - # Add device parameter for newer versions that support it - if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"): - kwargs["device"] = map_location - - # Load the model with appropriate kwargs - missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs) + missing_keys, unexpected_keys = load_model_as_safetensor( + model, model_file, strict=strict, device=map_location + ) log_model_loading_keys(missing_keys, unexpected_keys) - - # For older versions, manually move to device if needed - if "device" not in kwargs and map_location != "cpu": - logging.warning( - "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors." - " This means that the model is loaded on 'cpu' first and then copied to the device." - " This leads to a slower loading time." - " Please update safetensors to version 0.4.3 or above for improved performance." - ) - model.to(map_location) return model @abc.abstractmethod diff --git a/src/lerobot/rewards/pretrained.py b/src/lerobot/rewards/pretrained.py index d44b31733..ed4205bcc 100644 --- a/src/lerobot/rewards/pretrained.py +++ b/src/lerobot/rewards/pretrained.py @@ -21,8 +21,6 @@ from pathlib import Path from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any, TypeVar -import packaging -import safetensors from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError @@ -129,29 +127,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC): @classmethod def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T: - # Create base kwargs - kwargs = {"strict": strict} - - # Add device parameter for newer versions that support it - if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"): - kwargs["device"] = map_location - - # Load the model with appropriate kwargs - missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs) + missing_keys, unexpected_keys = load_model_as_safetensor( + model, model_file, strict=strict, device=map_location + ) if missing_keys: logging.warning(f"Missing key(s) when loading model: {missing_keys}") if unexpected_keys: logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}") - - # For older versions, manually move to device if needed - if "device" not in kwargs and map_location != "cpu": - logging.warning( - "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors." - " This means that the model is loaded on 'cpu' first and then copied to the device." - " This leads to a slower loading time." - " Please update safetensors to version 0.4.3 or above for improved performance." - ) - model.to(map_location) return model def get_optim_params(self):