Fix for codes not using pinned memory (apache#7001)

* Fix for codes not using pinned memory * Fix from review * Fix indentation
Guneet-Dhillon · Jul 11, 2017 · 2b9e960 · 2b9e960
1 parent 310633a
commit 2b9e960
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/storage/storage.cc b/src/storage/storage.cc
@@ -88,6 +88,8 @@ Storage::Handle StorageImpl::Alloc(size_t size, Context ctx) {
           }
           case Context::kGPU: {
 #if MXNET_USE_CUDA
+            CUDA_CALL(cudaGetDeviceCount(&num_gpu_device));
+            CHECK_GT(num_gpu_device, 0) << "GPU usage requires at least 1 GPU";
             ptr = new storage::GPUPooledStorageManager();
 #else
             LOG(FATAL) << "Compile with USE_CUDA=1 to enable GPU usage";