reference/generated/tensordict.make_tensordict

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

make_tensordict

tensordict.make_tensordict(input_dict: dict[str, Tensor | TensorCollection] | None = None, batch_size: Sequence[int] | Size | int | None = None, device: device | str | int | None = None, auto_batch_size: bool | None = None, **kwargs: Tensor | TensorCollection) → TensorDict

Returns a tensordict built from a dict or from keyword arguments.

Same as tensordict.TensorDict.from_dict(). Nested dicts become nested tensordicts. With auto_batch_size=True, each nested tensordict gets the largest batch size that its own entries share, which starts with the batch size of its parent.

Parameters:

• input_dict (dict or TensorDictBase, optional) – the entries. Its entries are added to those of **kwargs and replace them on a shared key. Defaults to None.

• batch_size (Size, list of int or int, optional) – the batch size of the result: the leading dimensions that all its entries share. Defaults to None, which gives [] unless auto_batch_size is True.

• device (device, str or int, optional) – the device of the result. Leaves on another device are moved to it. Defaults to None, which leaves each leaf where it is.

• auto_batch_size (bool, optional) – if True and batch_size is not given, sets the batch size to the largest leading shape that all the entries share. Defaults to None, which does not.

• **kwargs – the entries, as keyword arguments.

Examples

>>> import torch
>>> from tensordict import make_tensordict
>>> td = make_tensordict(a=torch.zeros(3, 4), b=torch.zeros(3))
>>> td.batch_size
torch.Size([])
>>> td = make_tensordict(
...     {"a": torch.zeros(3), "b": {"c": torch.zeros(3, 4)}}, auto_batch_size=True
... )
>>> td.batch_size, td["b"].batch_size
(torch.Size([3]), torch.Size([3, 4]))