fix(gpu feats)

graphistry · Sep 15, 2024 · f69cc8f · f69cc8f
1 parent 92dd7fa
commit f69cc8f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/graphistry/umap_utils.py b/graphistry/umap_utils.py
@@ -668,7 +668,10 @@ def umap(
             if isinstance(X_, pd.DataFrame):
                 index_to_nodes_dict = dict(zip(range(len(nodes)), nodes))
             elif 'cudf.core.dataframe' in str(getmodule(X_)):
-                index_to_nodes_dict = nodes  # {}?
+                assert isinstance(X_, cudf.DataFrame)
+                logger.debug('nodes type: %s', type(nodes))
+                import cupy as cp
+                index_to_nodes_dict = dict(zip(range(len(nodes)), cp.asnumpy(nodes)))
 
             # add the safe coercion here 
             X_, y_ = make_safe_gpu_dataframes(X_, y_, res.engine)  # type: ignore