reference/generated/tensordict.TensorDictBase.record_stream

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.record_stream

TensorDictBase.record_stream(stream: Stream) → Self

Marks the tensor leaves of the tensordict as used by stream.

Calls record_stream() on every tensor leaf: when a leaf is freed, its memory is not reused until the work queued on stream at that time is complete. A tensordict can be marked with one stream only.

Parameters:

stream (Stream) – the stream that uses the leaves.

Returns:

self.

Raises:

RuntimeError – if the tensordict is already marked with another stream.

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