diff --git a/python/cuml/cuml/internals/validation.py b/python/cuml/cuml/internals/validation.py index fe89807df6..2e00a2c623 100644 --- a/python/cuml/cuml/internals/validation.py +++ b/python/cuml/cuml/internals/validation.py @@ -683,14 +683,10 @@ def check_array( array = array.to_numpy(dtype="object") array = np.asarray(array, dtype=dtype, order=order) else: - # XXX: the dtype keyword to `to_cupy` is buggy, and also - # doesn't support all dtype coercions. For now we do a - # manual cast to handle any coercions. - # See https://github.com/rapidsai/cudf/issues/22136. - if dtype is not None: - array = array.astype(dtype, copy=False) array = cp.asarray( - array.to_cupy(copy=copy), dtype=dtype, order=order + array.to_cupy(copy=copy, dtype=dtype), + dtype=dtype, + order=order, ) elif isinstance(array, (pd.DataFrame, pd.Series)): # Handle pandas inputs diff --git a/python/cuml/tests/test_validation.py b/python/cuml/tests/test_validation.py index 029af195c9..7d67d1c603 100644 --- a/python/cuml/tests/test_validation.py +++ b/python/cuml/tests/test_validation.py @@ -1865,7 +1865,10 @@ def test_check_cudf(kind, ensure_ndim, coerce_ndim, dtype, shape): data = data.ravel() else: assert isinstance(out, (cudf.Series, cudf.DataFrame)) - res = out.to_numpy(dtype=data.dtype) + + res = out.to_numpy( + dtype=("float32" if data.dtype == "float16" else data.dtype) + ) np.testing.assert_array_equal(res, data)