Skip to content

Commit c777637

Browse files
committed
Accept memory resource (and stream) in copy
Changes `copy` to allow a memory resource to be specified (using the default device memory resource if unspecified). Also accepts a `stream` argument (similar to other functions).
1 parent 548b353 commit c777637

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

python/rmm/_lib/device_buffer.pyx

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,9 @@ cdef class DeviceBuffer:
146146
}
147147
return intf
148148

149-
def copy(self):
149+
def copy(self, *,
150+
Stream stream=DEFAULT_STREAM,
151+
DeviceMemoryResource mr=None):
150152
"""Returns a copy of DeviceBuffer.
151153
152154
Returns
@@ -165,9 +167,9 @@ cdef class DeviceBuffer:
165167
>>> assert db is not db_copy
166168
>>> assert db.ptr != db_copy.ptr
167169
"""
168-
ret = DeviceBuffer(ptr=self.ptr, size=self.size, stream=self.stream)
169-
ret.mr = self.mr
170-
return ret
170+
return DeviceBuffer(
171+
ptr=self.ptr, size=self.size, stream=stream, mr=mr
172+
)
171173

172174
def __copy__(self):
173175
return self.copy()

0 commit comments

Comments
 (0)