diff --git a/src/common/common.cu b/src/common/common.cu
index d51bd8c262db..43fa4d4652d0 100644
--- a/src/common/common.cu
+++ b/src/common/common.cu
@@ -11,7 +11,7 @@ int AllVisibleImpl::AllVisible() {
     // When compiled with CUDA but running on CPU only device,
     // cudaGetDeviceCount will fail.
     dh::safe_cuda(cudaGetDeviceCount(&n_visgpus));
-  } catch(const thrust::system::system_error& err) {
+  } catch(const dmlc::Error &except) {
     return 0;
   }
   return n_visgpus;
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 81abed25c4d7..507325e2271f 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -956,7 +956,7 @@ class SaveCudaContext {
     // cudaGetDevice will fail.
     try {
       safe_cuda(cudaGetDevice(&saved_device_));
-    } catch (const thrust::system::system_error & err) {
+    } catch (const dmlc::Error &except) {
       saved_device_ = -1;
     }
     func();