Handle cudaErrorNoDevice (#8184)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Oneflow-Inc · May 11, 2022 · 08e1e0b · 08e1e0b
1 parent d653f59
commit 08e1e0b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/oneflow/core/ep/cuda/cuda_device_manager.cpp b/oneflow/core/ep/cuda/cuda_device_manager.cpp
@@ -45,7 +45,9 @@ size_t CudaDeviceManager::GetDeviceCount(size_t primary_device_index) {
 
 size_t CudaDeviceManager::GetDeviceCount() {
   int count = 0;
-  OF_CUDA_CHECK(cudaGetDeviceCount(&count));
+  cudaError_t err = cudaGetDeviceCount(&count);
+  if (err == cudaErrorNoDevice) { return 0; }
+  OF_CUDA_CHECK(err);
   return count;
 }