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 the tensordict along a batch dimension, as torch.split() does.

The pieces hold views of the original tensors, and keep the names.

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 of int) – the size of each piece (the last one can be smaller), or the list of the sizes of the pieces, which must add up to batch_size[dim].

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

Returns:

A tuple of tensordicts.

Raises:

RuntimeError – if the sizes of a list do not add up to batch_size[dim].

Examples

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

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