Rmm
https://github.com/rapidsai/rmm Â
from rmm.allocators.torch import rmm_torch_allocator
from rmm.allocators.cupy import rmm_cupy_allocator
import torch
import cupy
torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
cupy.cuda.set_allocator(rmm_cupy_allocator)