TensorDictBase.size#

TensorDictBase.size(dim: int | None = None) → Size | int#

Returns the batch size, or the size of one batch dimension.

Parameters:: dim (int, optional) – a batch dimension. If None, the whole batch_size is returned. Defaults to None.

Examples

>>> from tensordict import TensorDict
>>> td = TensorDict(batch_size=[3, 4])
>>> td.size()
torch.Size([3, 4])
>>> td.size(-1)
4

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

Docs