Use `DeviceBuffer` from newer RMM releases
If a newer version of RMM is around, use `DeviceBuffer` instead.
Otherwise fallback to `device_array`. This is significantly faster to
allocate.
@@ -46,6 +46,9 @@ def init_once():
try:
import rmm
+ if hasattr(rmm, "DeviceBuffer"):
+ cuda_array = lambda n: rmm.DeviceBuffer(size=n)
+ else: # pre-0.12.0
cuda_array = lambda n: rmm.device_array(n, dtype=np.uint8)
except ImportError:
try:
