fix rmm reinitialize call

rapidsai · Nov 6, 2024 · a4b9f7e · a4b9f7e
1 parent 57bc48f
commit a4b9f7e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/cugraph-dgl/examples/dataset_from_disk_cudf.ipynb b/python/cugraph-dgl/examples/dataset_from_disk_cudf.ipynb
@@ -22,7 +22,7 @@
     "import rmm\n",
     "import torch\n",
     "from rmm.allocators.torch import rmm_torch_allocator\n",
-    "rmm.reinitialize(initial_pool_size=15e9)\n",
+    "rmm.reinitialize(pool_allocator=True, initial_pool_size=15e9)\n",
     "#Switch to async pool in case of memory issues due to fragmentation of the pool\n",
     "#rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource(initial_pool_size=15e9))\n",
     "torch.cuda.memory.change_current_allocator(rmm_torch_allocator)"
@@ -254,7 +254,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.8"
   },
   "vscode": {
    "interpreter": {