From 5c34d3dd340c76678d8f2667057c6b0ce2f1f480 Mon Sep 17 00:00:00 2001 From: Alex Barghi <105237337+alexbarghi-nv@users.noreply.github.com> Date: Tue, 26 Sep 2023 11:57:30 -0400 Subject: [PATCH] Update Allocator Selection in cuGraph-DGL Example (#3877) Closes #3847 Authors: - Alex Barghi (https://github.com/alexbarghi-nv) Approvers: - Vibhu Jawa (https://github.com/VibhuJawa) - Brad Rees (https://github.com/BradReesWork) URL: https://github.com/rapidsai/cugraph/pull/3877 --- .../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)