don't delete observation.image key

2024-05-08 12:58:14 +01:00 · 2024-05-08 12:58:14 +01:00 · c3a5cbb0b6
parent 8357dae26a
commit c3a5cbb0b6
1 changed files with 3 additions and 2 deletions
--- a/lerobot/common/policies/diffusion/modeling_diffusion.py
+++ b/lerobot/common/policies/diffusion/modeling_diffusion.py
@ -89,6 +89,7 @@ class DiffusionPolicy(nn.Module, PyTorchModelHubMixin):
            assert "action" in batch
            assert "action_is_pad" in batch
        image_key = next(iter(image_keys))
+        if image_key != "observation.image":
            batch["observation.image"] = batch[image_key]
            del batch[image_key]