Merge pull request #2921 from buaaliyi/multi_gpu

Destroy CUDA stream when finished
BVLC · Aug 14, 2015 · 65c7fa6 · 65c7fa6
2 parents 33e5b39 + e696f85
commit 65c7fa6
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/caffe/layers/base_data_layer.cpp b/src/caffe/layers/base_data_layer.cpp
@@ -74,7 +74,7 @@ void BasePrefetchingDataLayer<Dtype>::InternalThreadEntry() {
 #ifndef CPU_ONLY
   cudaStream_t stream;
   if (Caffe::mode() == Caffe::GPU) {
-    cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
+    CUDA_CHECK(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
   }
 #endif
 
@@ -85,14 +85,19 @@ void BasePrefetchingDataLayer<Dtype>::InternalThreadEntry() {
 #ifndef CPU_ONLY
       if (Caffe::mode() == Caffe::GPU) {
         batch->data_.data().get()->async_gpu_push(stream);
-        cudaStreamSynchronize(stream);
+        CUDA_CHECK(cudaStreamSynchronize(stream));
       }
 #endif
       prefetch_full_.push(batch);
     }
   } catch (boost::thread_interrupted&) {
     // Interrupted exception is expected on shutdown
   }
+#ifndef CPU_ONLY
+  if (Caffe::mode() == Caffe::GPU) {
+    CUDA_CHECK(cudaStreamDestroy(stream));
+  }
+#endif
 }
 
 template <typename Dtype>