reference/generated/tensordict.TensorDictBase.squeeze

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.squeeze

TensorDictBase.squeeze(*args, **kwargs)

Removes batch dimensions of size 1 from every tensor, and returns them in a new tensordict.

The tensors of the result are views of the original tensors. Feature dimensions of size 1 are kept.

Parameters:

dim (int | None) – the batch dimension to remove, in [-batch_dims, batch_dims - 1]. If dim is None, all singleton batch dimensions will be squeezed. Defaults to None.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> td = td.squeeze()
>>> td.shape
torch.Size([3, 4])
>>> td.get("x").shape
torch.Size([3, 4, 2])

This operation can be used as a context manager too. Inside the with block, you work on the squeezed tensordict. When the block exits, the inverse operation (unsqueeze()) runs, and the entries of the block are written back to the original tensordict. They replace the original tensors instead of being copied into them. The original tensordict must not be locked. If the block raises an exception, nothing is written back. Pass dim: with td.squeeze(): raises RuntimeError, because the inverse operation needs the dimension.

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> with td.squeeze(1) as tds:
...     _ = tds.set("y", torch.zeros(3, 4))
>>> assert td.get("y").shape == (3, 1, 4)

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

Tutorials that use squeeze

Changing the shape of a TensorDict

Changing the shape of a TensorDict