torch.Tensor.record_stream

Tensor.record_stream(stream)

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

Примечание

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