diff --git a/src/caffe/layers/base_data_layer.cpp b/src/caffe/layers/base_data_layer.cpp index 20f76f62994..5303fe9c5d3 100644 --- a/src/caffe/layers/base_data_layer.cpp +++ b/src/caffe/layers/base_data_layer.cpp @@ -74,7 +74,7 @@ void BasePrefetchingDataLayer::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::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::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