torch.cuda.change_current_allocator — PyTorch 2.7 documentation (original) (raw)

torch.cuda.change_current_allocator(allocator)[source][source]¶

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.