From c574eb49845d48f5aad532d823ef56aec1c0d0f2 Mon Sep 17 00:00:00 2001
From: Ilia Larchenko <41329713+IliaLarchenko@users.noreply.github.com>
Date: Fri, 14 Feb 2025 06:03:55 +0700
Subject: [PATCH] Fixed eval.py on MPS (#702)

---
 lerobot/scripts/eval.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py
index 7318748fb..a4f79afc4 100644
--- 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)