torch.utils.dlpack

torch.utils.dlpack.from_dlpack(ext_tensor) → Tensor

将外部库中的张量转换为一个 torch.Tensor 通过 __dlpack__ 协议实现。

该张量将与 DLPack 中表示的对象共享内存。

警告

每个胶囊只能调用一次 from_dlpack。对其在同一个胶囊上多次调用的行为未定义。

Parameters

ext_tensor (具有 __dlpack__ 属性的对象 或 DLPack 封装体) – 要转换的张量或 DLPack 封装体。

torch.utils.dlpack.to_dlpack(tensor) → PyCapsule

返回一个表示张量的 DLPack。

Parameters

tensor – 要导出的张量

DLPack 共享张量的内存。 请注意，每个 DLPack 只能被消费一次。