set_list_to_stack

class tensordict.set_list_to_stack(mode: bool)

Context manager and decorator to control the behavior of list handling in TensorDict.

When enabled, lists assigned to a TensorDict will be automatically stacked along the batch dimension. This can be useful for ensuring that lists of tensors or other elements are treated as stackable entities within a TensorDict.

Current Behavior:

If a list is assigned to a TensorDict without this context manager, it will be converted to a numpy array

and wrapped in a NonTensorData if it cannot be cast to a Tensor.

Future Behavior:

In version 0.10.0, lists will be automatically stacked by default.

Parameters:

mode (bool) – If True, enables list-to-stack conversion. If False, disables it.

Warning

A FutureWarning will be raised if a list is assigned to a TensorDict without setting this context manager

or the global flag, indicating that the behavior will change in the future.

Example

>>> with set_list_to_stack(True):
...     td = TensorDict(a=[torch.zeros(()), torch.ones(())], batch_size=2)
...     assert (td["a"] == torch.tensor([0, 1])).all()
...     assert td[0]["a"] == 0
...     assert td[1]["a"] == 1

See also

list_to_stack().