TensorDictBase.unsqueeze#
- TensorDictBase.unsqueeze(*args, **kwargs)#
Inserts a batch dimension of size 1 in every tensor, and returns them in a new tensordict.
The tensors of the result are views of the original tensors.
- Parameters:
dim (int) – the position of the new dimension, in
[-batch_dims - 1, batch_dims]as fortorch.unsqueeze().
Examples
>>> import torch >>> from tensordict import TensorDict >>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> td = td.unsqueeze(-2) >>> td.shape torch.Size([3, 1, 4]) >>> td.get("x").shape torch.Size([3, 1, 4, 2])
This operation can be used as a context manager too. Inside the
withblock, you work on the unsqueezed tensordict. When the block exits, the inverse operation (squeeze()) 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.>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> with td.unsqueeze(-2) as tds: ... _ = tds.set("y", torch.zeros(3, 1, 4)) >>> assert td.get("y").shape == (3, 4)
Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.
