Files
pytorch/torch/utils/dlpack.py
Thomas Viehmann b2a74d17ad document torch.utils.dlpack (#9343)
Summary:
dlpacks deserve documentation. :)

I wonder whether it might make sense to merge the various small torch.utils pages (and include a link for the larger ones, e.g. data) to enhance the structure in the docs.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/9343

Differential Revision: D8801227

Pulled By: soumith

fbshipit-source-id: 2980d271971743b86f052bec5a2cb4d146a90d9b
2018-07-11 07:46:09 -07:00

28 lines
656 B
Python

import torch
from torch._C import _from_dlpack as from_dlpack
from torch._C import _to_dlpack as to_dlpack
torch._C._add_docstr(from_dlpack, r"""from_dlpack(dlpack) -> Tensor
Decodes a DLPack to a tensor.
Arguments::
dlpack - a PyCapsule object with the dltensor
The tensor will share the memory with the object represented
in the dlpack.
Note that each dlpack can only be consumed once.
""")
torch._C._add_docstr(to_dlpack, r"""to_dlpack(tensor) -> PyCapsule
Returns a DLPack representing the tensor.
Arguments::
tensor - a tensor to be exported
The dlpack shares the tensors memory.
Note that each dlpack can only be consumed once.
""")