torch.Tensor.record_stream
Tensor.record_stream(stream)
Ensures that the tensor memory is not reused for another tensor
   until all current work queued on "stream" are complete.
Note:
 The caching allocator is aware of only the stream where a tensor
 was allocated. Due to the awareness, it already correctly manages
 the life cycle of tensors on only one stream. But if a tensor is
 used on a stream different from the stream of origin, the
 allocator might reuse the memory unexpectedly. Calling this
 method lets the allocator know which streams have used the
 tensor.
