当前位置: 首页 > 工具软件 > DLPack > 使用案例 >

Pytorch文档学习 TORCH.FROM_DLPACK

洪弘亮
2023-12-01

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

 类似资料: