dropped rmm from predict

EliHei2 · Dec 11, 2024 · f6b9ccd · f6b9ccd
1 parent 09ed69f
commit f6b9ccd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/segger/prediction/predict.py b/src/segger/prediction/predict.py
@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 import torch._dynamo
 import gc
-import rmm
+# import rmm
 import re
 import glob
 from pathlib import Path
@@ -28,7 +28,7 @@
 from dask.diagnostics import ProgressBar
 import time
 import dask
-from rmm.allocators.cupy import rmm_cupy_allocator
+# from rmm.allocators.cupy import rmm_cupy_allocator
 from cupyx.scipy.sparse import coo_matrix
 from torch.utils.dlpack import to_dlpack, from_dlpack