From 3921faa385503fa6ed4625494a1463232f4228bf Mon Sep 17 00:00:00 2001 From: eordentlich Date: Thu, 16 Jul 2026 20:45:41 -0700 Subject: [PATCH] assume cupy for sam when returning embeddings in umap (#1029) Signed-off-by: Erik Ordentlich --- python/src/spark_rapids_ml/umap.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/src/spark_rapids_ml/umap.py b/python/src/spark_rapids_ml/umap.py index 03cce339..e311cd46 100644 --- a/python/src/spark_rapids_ml/umap.py +++ b/python/src/spark_rapids_ml/umap.py @@ -1256,7 +1256,7 @@ def _train_udf(pdf_iter: Iterable[pd.DataFrame]) -> Iterable[pd.DataFrame]: indices = csr_chunk.indices indptr = csr_chunk.indptr data = csr_chunk.data - if cuda_managed_mem_enabled: + if cuda_managed_mem_enabled or cuda_system_mem_enabled: yield pd.DataFrame( data=[ { @@ -1281,7 +1281,7 @@ def _train_udf(pdf_iter: Iterable[pd.DataFrame]) -> Iterable[pd.DataFrame]: ] ) else: - if cuda_managed_mem_enabled: + if cuda_managed_mem_enabled or cuda_system_mem_enabled: yield pd.DataFrame( { "embedding_": list(embedding[start:end].get()),