From 4a2b3969e3fe7c512741348ba7e8e784e9cc8209 Mon Sep 17 00:00:00 2001 From: captainfff Date: Wed, 30 Jul 2025 08:41:30 +0000 Subject: [PATCH] fix pi0 prepare_language will raise an error if the task is a string --- src/lerobot/policies/pi0/modeling_pi0.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py index a34aa34f98..feff4bdc0d 100644 --- a/src/lerobot/policies/pi0/modeling_pi0.py +++ b/src/lerobot/policies/pi0/modeling_pi0.py @@ -392,6 +392,11 @@ def prepare_language(self, batch) -> tuple[Tensor, Tensor]: """Tokenize the text input""" device = batch[OBS_STATE].device tasks = batch["task"] + if isinstance(tasks, str): + tasks = [tasks] + + if len(tasks) == 1: + tasks = [tasks[0] for _ in range(batch[OBS_STATE].shape[0])] # PaliGemma prompt has to end with a new line tasks = [task if task.endswith("\n") else f"{task}\n" for task in tasks]