Fixed eval.py on MPS (#702)

huggingface · Feb 13, 2025 · c574eb4 · c574eb4
1 parent 1e49cc4
commit c574eb4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py
@@ -151,7 +151,9 @@ def rollout(
         if return_observations:
             all_observations.append(deepcopy(observation))
 
-        observation = {key: observation[key].to(device, non_blocking=True) for key in observation}
+        observation = {
+            key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
+        }
 
         with torch.inference_mode():
             action = policy.select_action(observation)