reference/generated/tensordict.TensorDictBase.set_

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.set

TensorDictBase.set_(key: NestedKey, item: Tensor | TensorCollection, *, non_blocking: bool = False) → Self

Copies a value into an existing entry, which keeps its storage.

The existing tensor keeps its data_ptr(), so other references to it see the new values. set_ works on a locked tensordict. See In-place operations.

Parameters:

• key (NestedKey) – the key of the entry. The key must exist.

• item (Tensor or compatible type, TensorDictBase) – value to be copied into the entry. Its shape must match the entry.

Keyword Arguments:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

Returns:

self

Raises:

• KeyError – if the key is missing.

• RuntimeError – if the shape of the value does not match the batch size.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({}, batch_size=[3, 4])
>>> x = torch.randn(3, 4)
>>> td = td.set("x", x)
>>> td = td.set_("x", torch.zeros_like(x))
>>> assert (x == 0).all()

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

Tutorials that use set_

Getting, setting, and nesting keys

Getting, setting, and nesting keys