reference/generated/tensordict.TensorDictBase.narrow

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.narrow

TensorDictBase.narrow(dim: int, start: int, length: int)

Returns a new tensordict that is a narrowed version of the input.

The dimension dim is input from start to start + length.

Parameters:

• dim (int) – The dimension along which to narrow.

• start (int) – Starting index.

• length (int) – Length of the narrowed dimension.

Returns:

a new tensordict narrowed along the specified dimension.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3])
>>> print(td["a"])
tensor([[0, 1, 2],
        [3, 4, 5]])
>>> print(td.narrow(1, 1, 2)["a"])
tensor([[1, 2],
        [4, 5]])

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