reference/generated/tensordict.TensorDictBase.chunk

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.chunk

TensorDictBase.chunk(chunks: int, dim: int = 0) → tuple[TensorCollection, ...]

Splits the tensordict into chunks pieces along a batch dimension, as torch.chunk() does.

Each piece has ceil(batch_size[dim] / chunks) elements, except the last one, which can be smaller. So there can be fewer than chunks pieces. The pieces hold views of the original tensors.

Parameters:

• chunks (int) – the number of pieces.

• dim (int, optional) – the batch dimension to split. Defaults to 0.

Returns:

A tuple of tensordicts.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(x=torch.arange(5), batch_size=[5])
>>> [t["x"] for t in td.chunk(2)]
[tensor([0, 1, 2]), tensor([3, 4])]
>>> [t["x"] for t in td.chunk(4)]
[tensor([0, 1]), tensor([2, 3]), tensor([4])]

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.