reference/generated/tensordict.TensorDictBase.lock_

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.lock

TensorDictBase.lock_() → Self

Locks the tensordict and its sub-tensordicts, and returns it.

A locked tensordict raises RuntimeError on any operation that adds, removes or renames an entry, or that replaces a leaf with a new tensor: set(), __setitem__, update(), pop(), rename_key_(), del. Writes into the existing leaves still work: set_(), update_(), and set(..., inplace=True). Check the state with is_locked. See Keys, nested keys, and locking.

Can be used as a context manager: the tensordict is unlocked when the with block exits, also if the block raises an exception.

Examples

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[])
>>> with td.lock_():
...     assert td.is_locked
...     try:
...         td.set("d", 0)  # adds a key: error
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         del td["a"]  # removes a key: error
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         td.rename_key_("a", "d")  # renames a key: error
...     except RuntimeError:
...         print("td is locked!")
...     # These calls write into the existing leaves, so they work:
...     _ = td.set("a", 0, inplace=True)
...     _ = td.set_("a", 0)
...     _ = td.update({"a": 0}, inplace=True)
...     _ = td.update_({"a": 0})
td is locked!
td is locked!
td is locked!
>>> td.is_locked
False

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

Tutorials that use lock_

Changing the shape of a TensorDict

Changing the shape of a TensorDict