cupy.cuda.using_allocator

cupy.cuda.using_allocator(allocator=None)

在上下文管理器内设置线程本地的 GPU 内存分配器

上下文管理器

参数:

allocator (function) – CuPy 内存分配器。它必须与 cupy.cuda.alloc() 函数具有相同的接口，该函数接受缓冲区大小作为参数并返回该大小的设备缓冲区。当指定 None 时，将使用原始内存分配器（即禁用内存池）。