diff options
| author | Ronghang Hu <huronghang@hotmail.com> | 2015-08-14 10:54:11 -0700 |
|---|---|---|
| committer | Ronghang Hu <huronghang@hotmail.com> | 2015-08-14 10:54:11 -0700 |
| commit | 65c7fa6e92d81c68d44caa206bd003013786ac69 (patch) | |
| tree | fc7b0cea4b15d16c4675aaad30e0ba5d83c3d7de /src | |
| parent | 33e5b3929cc2fa3fee8562408b7d4e8c0febb008 (diff) | |
| parent | e696f85abe2c50107fbb67b211bf7dad0f87ade0 (diff) | |
| download | caffeonacl-65c7fa6e92d81c68d44caa206bd003013786ac69.tar.gz caffeonacl-65c7fa6e92d81c68d44caa206bd003013786ac69.tar.bz2 caffeonacl-65c7fa6e92d81c68d44caa206bd003013786ac69.zip | |
Merge pull request #2921 from buaaliyi/multi_gpu
Destroy CUDA stream when finished
Diffstat (limited to 'src')
| -rw-r--r-- | src/caffe/layers/base_data_layer.cpp | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/src/caffe/layers/base_data_layer.cpp b/src/caffe/layers/base_data_layer.cpp index 20f76f62..5303fe9c 100644 --- 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,7 +85,7 @@ 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); @@ -93,6 +93,11 @@ void BasePrefetchingDataLayer<Dtype>::InternalThreadEntry() { } 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> |