reference/generated/tensordict.TensorDictBase.cat_tensors

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.cat_tensors

TensorDictBase.cat_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → Self

Concatenates entries of the tensordict into a new entry, and removes them by default.

Parameters:

*keys (NestedKey) – the entries to concatenate, with torch.cat().

Keyword Arguments:

• out_key (NestedKey) – the key of the result.

• dim (int, optional) – the dimension to concatenate along, counted in the dimensions of the entries. Defaults to 0.

• keep_entries (bool, optional) – if False, removes the entries of keys once the result is written. If the concatenation or the write fails, every entry is kept with its original value. Defaults to False.

Returns:

self.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1))
>>> _ = td.cat_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

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