reference/generated/tensordict.TensorDictBase.new_ones

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.new_ones

TensorDictBase.new_ones(*size: Size, dtype: dtype = None, device: device | str | int = NO_DEFAULT, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: bool | None = None, empty_lazy: bool = False)

Returns a tensordict of batch size size with every leaf filled with 1.

Each leaf keeps its feature dimensions: a leaf of shape (*self.batch_size, *f) becomes (*size, *f). By default, each leaf keeps its dtype and device.

Parameters:

size (int...) – the batch size of the result, as integers or as one list, tuple or torch.Size.

Keyword Arguments:

• dtype (dtype, optional) – the dtype of every leaf of the result. Defaults to None: each leaf keeps its dtype.

• device (device, optional) – the device of every leaf of the result. If not passed, each leaf keeps its device.

• requires_grad (bool, optional) – if True, autograd records the operations on the new leaves. Defaults to False.

• layout (layout, optional) – the layout of the new leaves. Defaults to torch.strided.

• pin_memory (bool, optional) – if True, the new leaves are allocated in pinned memory. Works only for CPU tensors. Defaults to False.

• empty_lazy (bool, optional) – if True, the lazy stacks of the result are emptied of their content. Use it when the content of a lazy stack is likely to change while you fill the new tensordict. It applies to the sub-tensordicts too. Defaults to False.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3, 4), b=torch.zeros(3, dtype=torch.int64), batch_size=[3])
>>> out = td.new_ones(2)
>>> out.batch_size, out["a"].shape
(torch.Size([2]), torch.Size([2, 4]))

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