def setup_rmm_pool(client, pool_size):
# Initialize an RMM pool allocator.
# Note: RMM may require the pool size to be a multiple of 256.
pool_size = get_rmm_size(pool_size)
client.run(rmm.reinitialize,
pool_allocator=True,
initial_pool_size=pool_size)
return None
def setup_rmm_pool(client, pool_size):
pool_size = get_rmm_size(pool_size)
client.run(rmm.reinitialize,
pool_allocator=True,
initial_pool_size=pool_size)
return None