torch.Tensor.record_stream

Tensor.record_stream(stream)

Гарантирует, что память тензора не будет повторно использована для другого тензора до тех пор, пока все текущие задачи, помещенные в очередь на stream, не будут завершены.

Примечание

Кэширующий выделенный блок памяти знает только о потоке, в котором был выделен тензор. Благодаря этому он уже правильно управляет жизненным циклом тензоров только в одном потоке. Но если тензор используется в потоке, отличном от потока происхождения, выделенный блок памяти может неожиданно повторно использовать память. Вызов этого метода позволяет выделенному блоку памяти узнать, какие потоки использовали тензор.