use mapper.dtype

rapidsai · Oct 20, 2023 · 6d43317 · 6d43317
1 parent 786c1e2
commit 6d43317
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/core.py b/python/nx-cugraph/nx_cugraph/algorithms/core.py
@@ -63,7 +63,7 @@ def k_truss(G, k):
         edge_masks = {key: val[edge_indices] for key, val in G.edge_masks.items()}
     # Renumber step 2: edge indices
     mapper = cp.zeros(len(G), src_indices.dtype)
-    mapper[node_indices] = cp.arange(node_indices.size, dtype=np.int64)
+    mapper[node_indices] = cp.arange(node_indices.size, dtype=mapper.dtype)
     src_indices = mapper[src_indices]
     dst_indices = mapper[dst_indices]
     # Renumber step 3: node values