diff --git a/dask_cuda/initialize.py b/dask_cuda/initialize.py
index 535c2b37c..8f159cab7 100644
--- 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__)