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, and returns self.

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, TensorDictBase or any compatible value) – the value to copy into the entry. Its leading dimensions must match the batch size, and it is broadcast to the shape of the entry.

Keyword Arguments:

non_blocking (bool, optional) – if True, copies between devices can run asynchronously with respect to the host. Defaults to False.

Returns:

self.

Raises:

• KeyError – if the key is missing.

• RuntimeError – if the leading dimensions of the value do not match the batch size.

• ValueError – if the value does not fit the shape of the entry. The error is also a RuntimeError, the type that update_() raises.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> x = torch.ones(3, 4)
>>> td = TensorDict(x=x, batch_size=[3, 4])
>>> td = td.set_("x", torch.zeros(3, 4))
>>> td["x"] is x, x.sum()
(True, tensor(0.))

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