Fix MultiGPU hangs when interrupting training

743341 · Aug 5, 2017 · feb5038 · feb5038
1 parent 7ad1342
commit feb5038
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/caffe/parallel.cpp b/src/caffe/parallel.cpp
@@ -354,11 +354,12 @@ void NCCL<Dtype>::Run(const vector<int>& gpus, const char* restore) {
   // Run first solver on current thread
   Broadcast();
   solver_->Solve();
-  barrier.wait();  // Hangs without it when running tests
   // Wait for shutdown
   for (int i = 1; i < gpus.size(); ++i) {
     workers[i]->StopInternalThread();
   }
+  // Hangs without it when running tests or multi-gpu training.
+  barrier.wait();
 }
 
 INSTANTIATE_CLASS(Params);