From 5524a0d7a7741bb1cfedc6a53b2298c8066daa69 Mon Sep 17 00:00:00 2001 From: Bryson Jones Date: Wed, 10 Dec 2025 13:50:47 -0800 Subject: [PATCH] split up select action return statement --- src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py b/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py index ec9138bf7..3620ad8fd 100644 --- a/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py +++ b/src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py @@ -795,7 +795,8 @@ class MultiTaskDiTPolicy(PreTrainedPolicy): actions = self.predict_action_chunk(batch) self._queues[ACTION].extend(actions.transpose(0, 1)) - return self._queues[ACTION].popleft() + action = self._queues[ACTION].popleft() + return action def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]: """Run the batch through the model and compute the loss for training"""