Merge branch 'branch-23.10' into wholegraph-fs

rapidsai · Sep 26, 2023 · 3dee752 · 3dee752
2 parents 9967da4 + 5c34d3d
commit 3dee752
Showing 1 changed file with 5 additions and 3 deletions.
diff --git 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)