cupy.cuda.using_allocator

cupy.cuda.using_allocator(allocator=None)

Sets a thread-local allocator for GPU memory inside

context manager

Parameters:

allocator (function) – CuPy memory allocator. It must have the same interface as the cupy.cuda.alloc() function, which takes the buffer size as an argument and returns the device buffer of that size. When None is specified, raw memory allocator will be used (i.e., memory pool is disabled).