From ae7a54de578a2b8bb81c830c0e8d0a8946aa38ab Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Tue, 22 Jul 2025 11:03:28 +0200 Subject: [PATCH] refactor(pipeline): Utilize get_safe_torch_device for device assignment - Replaced direct torch.device instantiation with get_safe_torch_device to ensure safe device handling. - This change enhances code readability and maintains consistency in device management across the RobotProcessor class. --- src/lerobot/processor/pipeline.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index 888ec3b36..ea08e2d1a 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -28,6 +28,8 @@ import torch from huggingface_hub import ModelHubMixin, hf_hub_download from safetensors.torch import load_file, save_file +from lerobot.utils.utils import get_safe_torch_device + class TransitionKey(str, Enum): """Keys for accessing EnvTransition dictionary components.""" @@ -465,7 +467,7 @@ class RobotProcessor(ModelHubMixin): to the target device, and reload it. Only works for steps that implement both state_dict() and load_state_dict() methods. """ - device = torch.device(device) + device = get_safe_torch_device(device) for step in self.steps: if hasattr(step, "state_dict") and hasattr(step, "load_state_dict"):