torch.cuda.change_current_allocator
torch.cuda.change_current_allocator(allocator)
Changes 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
