reference/generated/tensordict.TensorDictBase.swapdims

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.swapdims

TensorDictBase.swapdims(axis0: int, axis1: int)

Returns a tensordict with two batch dimensions swapped, as transpose() does.

swapaxes and swapdims are aliases of transpose(). It can be used as a context manager, as transpose() can.

Parameters:

• axis0 (int) – the first batch dimension to swap.

• axis1 (int) – the second batch dimension to swap.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.zeros(3, 4, 5)}, batch_size=[3, 4])
>>> print(td.swapaxes(0, 1).shape)
torch.Size([4, 3])

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