MCPcopy Create free account
hub / github.com/CVCUDA/CV-CUDA / test_wrap_torch_buffer_dlpack

Function test_wrap_torch_buffer_dlpack

tests/cvcuda/python/test_tensor.py:112–129  ·  view source on GitHub ↗
(shape, dtype)

Source from the content-addressed store, hash-verified

110
111@t.mark.parametrize("shape,dtype", params_wrap_torch)
112def test_wrap_torch_buffer_dlpack(shape, dtype):
113 ttensor = torch.as_tensor(np.ndarray(shape, dtype=dtype), device="cuda")
114
115 # Since cvcuda.as_tensor can understand both dlpack and cuda_array_interface,
116 # and we don't know a priori which interfaces it'll use (torch provides both),
117 # let's create one object with only the dlpack interface.
118 class DLPackObject:
119 pass
120
121 o = DLPackObject()
122 o.__dlpack__ = ttensor.__dlpack__
123 o.__dlpack_device__ = ttensor.__dlpack_device__
124
125 tensor = cvcuda.as_tensor(o)
126 assert tensor.shape == shape
127 assert tensor.dtype == dtype
128 assert tensor.layout is None
129 assert tensor.ndim == len(shape)
130
131
132@t.mark.parametrize("shape,dtype", params_wrap_torch)

Callers

nothing calls this directly

Calls 1

DLPackObjectClass · 0.85

Tested by

no test coverage detected