reference/generated/tensordict.TensorDictBase.split

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.split

TensorDictBase.split(split_size: int | list[int], dim: int = 0) → tuple[TensorDictBase, ...]

Splits each tensor in the TensorDict with the specified size in the given dimension, like torch.split.

Returns a list of TensorDict instances with the view of split chunks of items.

On a LazyStackedTensorDict, sizes along the stack dim that are negative or do not add up to the stack size emit a FutureWarning, and the pieces miss members or have a wrong batch size. Starting with TensorDict 0.17, the default will change to raising a RuntimeError, as TensorDict.split does.

Parameters:

• split_size (int or List(int)) – size of a single chunk or list of sizes for each chunk.

• dim (int) – dimension along which to split the tensor.

Returns:

A list of TensorDict with specified size in given dimension.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1 = td.split([1, 2], dim=0)
>>> print(td0['x'])
tensor([[0, 1, 2, 3]])

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

Tutorials that use split

Changing the shape of a TensorDict

Changing the shape of a TensorDict