A dictionary that treats a pile of tensors like one tensor

What it does

TensorDict is a dictionary-like container built specifically for PyTorch. It holds heterogeneous tensors under named keys but exposes the familiar tensor API—indexing, slicing, reshaping, and device casting—applying each operation to every leaf automatically. A companion @tensorclass decorator offers a typed dataclass variant for tensor bundles, aiming to replace repetitive per-tensor boilerplate with single calls over an entire structure.

The interesting bit

The value lies in the structural assumption that your dictionary has a shared batch dimension. Because TensorDict treats the collection as a single tensor-like entity, you can write generic training loops where data loading, model prediction, and loss computation are swappable components—the same loop works for classification, segmentation, or RL. It also pushes device transfers asynchronously and supports torch.compile, which turns a convenience layer into a performance tool.

Key highlights

• Dispatches to("cuda"), stack, cat, and slicing across all contained tensors in one call.
• Supports nesting, lazy preallocation, and memory-mapped serialization for checkpointing.
• Compatible with torch.vmap and distributed training across workers and machines.
• @tensorclass provides a dataclass interface for typed tensor structures.
• Already integrated into downstream projects like TorchRL, MuJoCo Playground, and verl.

Verdict

If you manage batches of structured tensor data—especially in RL, robotics, or multi-modal pipelines—this removes a surprising amount of friction. If you only ever handle a single tensor or simple tuples, it is probably overkill.