Skip to content

[python] Avoid an extra shared-memory copy for LeRobot video batches - #10733

Merged
JingsongLi merged 2 commits into
apache:masterfrom
XiaoHongbo-Hope:python-lerobot-direct-shared-batch
Oct 22, 2026
Merged

JingsongLi merged 2 commits into
apache:masterfrom
XiaoHongbo-Hope:python-lerobot-direct-shared-batch

Conversation

@XiaoHongbo-Hope

Copy link
Copy Markdown
Contributor

Purpose

LeRobot video batches currently allocate a CPU tensor and then call share_memory_(), copying the entire uninitialized storage before filling it. Allocate shared storage directly, following PyTorch DataLoader's collation pattern, to avoid that copy.

The change applies to shared CPU batch output; other allocation paths are unchanged. No standalone throughput improvement is claimed.

Tests

  • Video parallelism suite: 14 passed, 26 subtests passed. The regression checks shared, contiguous uint8/float outputs, repeated/unsorted windows, sample isolation, and rejects calls to share_memory_().
  • LeRobot suite: 67 passed, 29 optional tests skipped, 72 subtests passed (PyTorch 2.7.1+cpu).
  • Repository flake8, Python 3.6 syntax compilation, and git diff --check passed.

@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as ready for review October 21, 2026 17:17
@JingsongLi

Copy link
Copy Markdown
Contributor

+1

@JingsongLi
JingsongLi merged commit 729b62e into apache:master Oct 22, 2026
8 of 18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants