reference/generated/tensordict.TensorDictBase.cat

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.cat

classmethod TensorDictBase.cat(input, dim: int = 0, *, out=None)

Concatenates tensordicts into a single tensordict along the given dimension.

This call is equivalent to calling torch.cat() but is compatible with torch.compile. The tensordicts must have the same keys, and the same batch size except along dim.

Parameters:

• input (sequence of TensorDictBase) – the tensordicts to concatenate.

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

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the result into.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(x=torch.ones(3), batch_size=[3])
>>> TensorDict.cat([a, b]).batch_size
torch.Size([6])

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