From a23008ea852da79da2aab972cb4940f8f72d53af Mon Sep 17 00:00:00 2001 From: ChengjieLi Date: Thu, 6 Jul 2023 21:33:19 +0800 Subject: [PATCH] use rmm.DeviceBuffer --- python/xoscar/backends/communication/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/xoscar/backends/communication/utils.py b/python/xoscar/backends/communication/utils.py index d4bc5509..df161c2a 100644 --- a/python/xoscar/backends/communication/utils.py +++ b/python/xoscar/backends/communication/utils.py @@ -74,8 +74,8 @@ async def read_buffers(header: Dict, reader: StreamReader): for is_cuda_buffer, buf_size in zip(is_cuda_buffers, buffer_sizes): if is_cuda_buffer: # pragma: no cover if buf_size == 0: - content = await reader.readexactly(buf_size) - buffers.append(content) + # uniformly use rmm.DeviceBuffer for cuda's deserialization + buffers.append(rmm.DeviceBuffer(size=buf_size)) else: buffer = rmm.DeviceBuffer(size=buf_size) arr = _convert_to_cupy_ndarray(buffer)