Update train.py

2025-02-22 18:11:06 +01:00 · 2025-02-22 18:11:06 +01:00 · aa70e14033
parent 6ca03b0dac
commit aa70e14033
1 changed files with 1 additions and 1 deletions
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@ -81,7 +81,7 @@ def custom_collate_fn(batch):
                )

                # stack frames for this video key and add directly to the item
-                item[vid_key] = torch.stack(frames)
+                item[vid_key] = frames

        # add item data (both video and non-video) to final_batch
        for key, value in item.items():