Add back has_cuda_context before next Distributed release

rapidsai · Sep 14, 2021 · 488ef3d · 488ef3d
1 parent e45cc34
commit 488ef3d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/dask_cuda/initialize.py b/dask_cuda/initialize.py
@@ -7,9 +7,8 @@
 
 import dask
 import distributed.comm.ucx
-from distributed.diagnostics.nvml import has_cuda_context
 
-from .utils import get_ucx_config
+from .utils import get_ucx_config, has_cuda_context
 
 logger = logging.getLogger(__name__)