[core] cudagraph output with tensor weak reference (#9724)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-10-27 00:19:28 -07:00
committed by GitHub
parent 67a6882da4
commit 8549c82660
4 changed files with 50 additions and 28 deletions

View File

@@ -1479,3 +1479,12 @@ class LazyDict(Mapping, Generic[T]):
def __len__(self):
return len(self._factory)
def weak_ref_tensor(tensor: torch.Tensor) -> torch.Tensor:
"""
Create a weak reference to a tensor.
The new tensor will share the same data as the original tensor,
but will not keep the original tensor alive.
"""
return torch.ops._C.weak_ref_tensor(tensor)