Fix _extract_partitions function.

rapidsai · Sep 12, 2023 · ed1040d · ed1040d
1 parent b4c34ef
commit ed1040d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/cuml/dask/common/part_utils.py b/python/cuml/dask/common/part_utils.py
@@ -160,7 +160,7 @@ def _extract_partitions(dask_obj, client=None):
 
     yield wait(parts)
 
-    key_to_part = [(str(part.key), part) for part in parts]
+    key_to_part = [(part.key, part) for part in parts]
     who_has = yield client.who_has(parts)
 
     raise gen.Return(