Python get_rmm_size Examples

Example #1

File: dask-nvtabular-criteo-benchmark.py Project: bschifferer/NVTabular

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

Example #2

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