diff --git a/src/caffe/layers/base_data_layer.cu b/src/caffe/layers/base_data_layer.cu index 56439bc506a..ff6e412aba6 100644 --- a/src/caffe/layers/base_data_layer.cu +++ b/src/caffe/layers/base_data_layer.cu @@ -20,7 +20,9 @@ void BasePrefetchingDataLayer::Forward_gpu( caffe_copy(batch->label_.count(), batch->label_.gpu_data(), top[1]->mutable_gpu_data()); } - + // Ensure the copy is synchronous wrt the host, so that the next batch isn't + // copied in meanwhile. + CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault)); prefetch_free_.push(batch); }