data layer race condition bugfix

diff --git a/examples/demo_mnist.cpp b/examples/demo_mnist.cpp
index 9eb6204..11d3fc5 100644 (file)
--- a/examples/demo_mnist.cpp
+++ b/examples/demo_mnist.cpp
@@ -25,6 +25,7 @@ int main(int argc, char** argv) {
     return 0;
   }
   google::InitGoogleLogging(argv[0]);
+  Caffe::DeviceQuery();
 
   if (argc == 4 && strcmp(argv[3], "GPU") == 0) {
     LOG(ERROR) << "Using GPU";

diff --git a/src/caffe/layers/data_layer.cpp b/src/caffe/layers/data_layer.cpp
index fd26c5d..9ed9516 100644 (file)
--- a/src/caffe/layers/data_layer.cpp
+++ b/src/caffe/layers/data_layer.cpp
@@ -170,8 +170,9 @@ void DataLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,
   // cpu_data calls so that the prefetch thread does not accidentally make
   // simultaneous cudaMalloc calls when the main thread is running. In some
   // GPUs this seems to cause failures if we do not so.
-  layer->prefetch_data_->mutable_cpu_data();
-  layer->prefetch_label_->mutable_cpu_data();
+  prefetch_data_->mutable_cpu_data();
+  prefetch_label_->mutable_cpu_data();
+  data_mean_.cpu_data();
   // LOG(INFO) << "Initializing prefetch";
   CHECK(!pthread_create(&thread_, NULL, DataLayerPrefetch<Dtype>,
       reinterpret_cast<void*>(this))) << "Pthread execution failed.";