reference/generated/tensordict.TensorDictBase.load

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.load

classmethod TensorDictBase.load(prefix: str | Path, *args, **kwargs) → Self

Loads a tensordict that save() or memmap() wrote to disk.

load calls load_memmap(), which documents the other arguments, such as device, mode, subpath and allow_pickle. The leaves of the result read the files lazily. See Save and load tensordicts.

Parameters:

• prefix (str or Path) – the directory, or the archive file, to read.

• *args – passed to load_memmap().

• **kwargs – passed to load_memmap().

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.arange(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.save(path)
...     loaded = TensorDict.load(path)
...     print(loaded["a"], loaded.batch_size)
MemoryMappedTensor([0, 1, 2]) torch.Size([3])

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