reference/generated/tensordict.utils.TensorDictFuture

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictFuture

class tensordict.utils.TensorDictFuture(futures, resulting_td)

Bases: object

A tensordict that is still being written by background threads.

memmap(), memmap_(), memmap_like(), save() and consolidate() return one when called with return_early=True and num_threads > 1. Call result() to wait for the writes and get the tensordict. You do not build one yourself.

Parameters:

• futures (list of Future) – the writes to wait for.

• resulting_td (TensorDictBase) – the tensordict that result() returns.

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), b=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as path:
...     future = td.memmap(path, num_threads=2, return_early=True)
...     result = future.result()
...     print(result.is_memmap(), result["b"].tolist())
True [1.0, 1.0, 1.0]

result()

Waits for the writes to finish and returns the tensordict.

If a write failed, raises the first error of the writes.