diff --git a/python/paddle/utils/dlpack.py b/python/paddle/utils/dlpack.py index d3a7e7918f765d96f26aaa091527ab106b15f2ea..ca2a1ae0e19ec56c1e4d64022aed6992996d819f 100644 --- a/python/paddle/utils/dlpack.py +++ b/python/paddle/utils/dlpack.py @@ -17,6 +17,11 @@ from ..fluid.core import LoDTensor from ..fluid.framework import in_dygraph_mode from ..fluid.data_feeder import check_type, check_dtype, convert_dtype +__all__ = [ + 'to_dlpack', + 'from_dlpack', +] + def to_dlpack(x): """ @@ -63,7 +68,8 @@ def to_dlpack(x): def from_dlpack(dlpack): - """Decodes a DLPack to a tensor. + """ + Decodes a DLPack to a tensor. Args: dlpack (PyCapsule): a PyCapsule object with the dltensor. @@ -82,8 +88,8 @@ def from_dlpack(dlpack): x = paddle.utils.dlpack.from_dlpack(dlpack) print(x) # Tensor(shape=[2, 4], dtype=float32, place=CUDAPlace(0), stop_gradient=True, - [[0.20000000, 0.30000001, 0.50000000, 0.89999998], - [0.10000000, 0.20000000, 0.60000002, 0.69999999]]) + # [[0.20000000, 0.30000001, 0.50000000, 0.89999998], + # [0.10000000, 0.20000000, 0.60000002, 0.69999999]]) """ t = type(dlpack)