From 6633c75796dd5f05f9276e4bb4bfab168faf212b Mon Sep 17 00:00:00 2001 From: Alexandria Barghi Date: Mon, 25 Sep 2023 13:44:29 -0700 Subject: [PATCH] update allocator set method --- .../cugraph-dgl/examples/graphsage/node-classification.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/cugraph-dgl/examples/graphsage/node-classification.py b/python/cugraph-dgl/examples/graphsage/node-classification.py index 24df73ada75..320890b0312 100644 --- a/python/cugraph-dgl/examples/graphsage/node-classification.py +++ b/python/cugraph-dgl/examples/graphsage/node-classification.py @@ -39,14 +39,16 @@ def set_allocators(): + import rmm import cudf import cupy - import rmm + from rmm.allocators.torch import rmm_torch_allocator + from rmm.allocators.cupy import rmm_cupy_allocator mr = rmm.mr.CudaAsyncMemoryResource() rmm.mr.set_current_device_resource(mr) - torch.cuda.memory.change_current_allocator(rmm.rmm_torch_allocator) - cupy.cuda.set_allocator(rmm.allocators.cupy.rmm_cupy_allocator) + torch.cuda.memory.change_current_allocator(rmm_torch_allocator) + cupy.cuda.set_allocator(rmm_cupy_allocator) cudf.set_option("spill", True)