Rate this Page
★ ★ ★ ★ ★

TensorDictBase.new_zeros#

TensorDictBase.new_zeros(*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 0.

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_zeros(2, 2)
>>> out.batch_size, out["a"].shape, out["b"].dtype
(torch.Size([2, 2]), torch.Size([2, 2, 4]), torch.int64)

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