check for nvidia driver's sufficiency before checking for number of C…

…UDA devices (#1156)
pytorch · Mar 31, 2017 · b13b701 · b13b701
1 parent 5c79046
commit b13b701
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/torch/csrc/autograd/engine.cpp b/torch/csrc/autograd/engine.cpp
@@ -350,7 +350,14 @@ auto Engine::ready_queue(int device) -> ReadyQueue& {
 auto Engine::start_threads() -> void {
   int num_devices = 0;
 #ifdef WITH_CUDA
-  THCudaCheck(cudaGetDeviceCount(&num_devices));
+  cudaError_t err = cudaGetDeviceCount(&num_devices);
+
+  // check for case of compiled with CUDA but no NVIDIA driver available
+  if (err == cudaErrorInsufficientDriver) {
+    num_devices = 0;
+  } else {
+    THCudaCheck(err);
+  }
 #endif
   ready_queues = std::vector<std::unique_ptr<ReadyQueue>>(num_devices + 1);
   for (auto& queue : ready_queues) {