diff --git a/lerobot/scripts/server/network_utils.py b/lerobot/scripts/server/network_utils.py index b24c8f17..c62f6cbd 100644 --- a/lerobot/scripts/server/network_utils.py +++ b/lerobot/scripts/server/network_utils.py @@ -111,7 +111,7 @@ def state_to_bytes(state_dict: dict[str, torch.Tensor]) -> bytes: def bytes_to_state_dict(buffer: bytes) -> dict[str, torch.Tensor]: buffer = io.BytesIO(buffer) buffer.seek(0) - return torch.load(buffer) # nosec B614: Safe usage of torch.load + return torch.load(buffer, weights_only=False) # nosec B614: Safe usage of torch.load def python_object_to_bytes(python_object: Any) -> bytes: @@ -129,7 +129,7 @@ def bytes_to_python_object(buffer: bytes) -> Any: def bytes_to_transitions(buffer: bytes) -> list[Transition]: buffer = io.BytesIO(buffer) buffer.seek(0) - transitions = torch.load(buffer) # nosec B614: Safe usage of torch.load + transitions = torch.load(buffer, weights_only=False) # nosec B614: Safe usage of torch.load # Add validation checks here return transitions