Merge branch '2312-dbg-dask-timeout' of github.com:dantegd/cuml into …

…2312-dbg-dask-timeout
rapidsai · Nov 29, 2023 · cc82a80 · cc82a80
2 parents 6146cd3 + 0592e7b
commit cc82a80
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/cuml/tests/dask/test_dask_nearest_neighbors.py b/python/cuml/tests/dask/test_dask_nearest_neighbors.py
@@ -30,12 +30,12 @@
 pd = cpu_only_import("pandas")
 
 np = cpu_only_import("numpy")
-cp = cpu_only_import("cupy")
+cp = gpu_only_import("cupy")
 
 
 IS_ARM = platform.processor() == "aarch64"
 
-if IS_ARM and cp.cuda.runtime.runtimeGetVersion() <= 11080:
+if IS_ARM and cp.cuda.runtime.runtimeGetVersion() < 11080:
     pytest.skip(
         "Test hang in AARCH64 with CUDA < 11.8: "
         "https://github.com/rapidsai/cuml/issues/5673",