torch.from_dlpack
torch.from_dlpack(ext_tensor) → Tensor
把一个额外库的张量转换为torch中的张量。
返回的 PyTorch 张量将与输入张量共享内存,因此要注意不能随意更改该张量中的内容。
import torch.utils.dlpack
t = torch.arange(4)
t2 = torch.from_dlpack(t)
t2[:2] = -1 # show that memory is shared
t2
t
capsule = torch.utils.dlpack.to_dlpack(t)
capsule
t3 = torch.from_dlpack(capsule)
t3
t3[0] = -9 # now we're sharing memory between 3 tensors
t3
t2
t