From 8ef1922abda3260fa2c6794f004af44dd6e6433e Mon Sep 17 00:00:00 2001 From: demobo-com Date: Mon, 9 Jun 2025 16:02:18 -0700 Subject: [PATCH] perform pyforch conversion using GPU --- lerobot/common/robot_devices/control_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index f8a2ffae90..a7be59938c 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -109,11 +109,11 @@ def predict_action(observation, policy, device, use_amp): ): # Convert to pytorch format: channel first and float32 in [0,1] with batch dimension for name in observation: + observation[name] = observation[name].to(device) if "image" in name: observation[name] = observation[name].type(torch.float32) / 255 observation[name] = observation[name].permute(2, 0, 1).contiguous() observation[name] = observation[name].unsqueeze(0) - observation[name] = observation[name].to(device) # Compute the next action with the policy # based on the current observation