torch.cuda.change_current_allocator

torch.cuda.change_current_allocator(allocator)

Change the currently used memory allocator to be the one provided.

If the current allocator has already been used/initialized, this function will error.

Parameters

allocator (torch.cuda.memory._CUDAAllocator) – allocator to be set as the active one.

Note

See Memory management for details on creating and using a custom allocator