diff --git a/python/cuml/dask/common/part_utils.py b/python/cuml/dask/common/part_utils.py index 39466bf075..a6aa892a76 100644 --- 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(