From 1df19ae4687bf6dbd48f98b5b541850f70b65e0a Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Wed, 1 Jul 2026 17:45:16 +0000 Subject: [PATCH] only call .detach().cpu() once per caemra instead of once per image --- src/lerobot/policies/evo1/modeling_evo1.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/evo1/modeling_evo1.py b/src/lerobot/policies/evo1/modeling_evo1.py index 3606d63ca..1d0de98de 100644 --- a/src/lerobot/policies/evo1/modeling_evo1.py +++ b/src/lerobot/policies/evo1/modeling_evo1.py @@ -342,10 +342,15 @@ class EVO1Policy(PreTrainedPolicy): image_batches: list[list[Tensor]] = [] image_masks = torch.zeros(batch_size, self.config.max_views, dtype=torch.bool) + cpu_images: dict[str, Tensor] = { + camera_key: normalized[camera_key].detach().cpu() + for camera_key in camera_keys[: self.config.max_views] + } + for batch_index in range(batch_size): sample_images: list[Tensor] = [] for camera_key in camera_keys[: self.config.max_views]: - sample_images.append(normalized[camera_key][batch_index].detach().cpu()) + sample_images.append(cpu_images[camera_key][batch_index]) if not sample_images: raise ValueError("EVO1 received a batch without any image tensor.") while len(sample_images) < self.config.max_views: