From 631aa3d10a33a1fbb52f9c6ec0ebd5022b80ede7 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Tue, 20 Mar 2018 15:38:26 +0800
Subject: [PATCH] Wait all inputs ready

---
 paddle/fluid/framework/parallel_executor.cc | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/paddle/fluid/framework/parallel_executor.cc b/paddle/fluid/framework/parallel_executor.cc
index 21a19cb5b27..248a1b4a257 100644
--- a/paddle/fluid/framework/parallel_executor.cc
+++ b/paddle/fluid/framework/parallel_executor.cc
@@ -375,6 +375,12 @@ struct NCCLAllReduceOpHandle : public OpHandle {
     if (this->inputs_.size() == 1) {
       return;  // No need to all reduce when GPU count = 1;
     } else {
+      // Wait input done
+      for (auto *in : inputs_) {
+        auto &p = static_cast<VarHandle *>(in)->place_;
+        in->generated_op_->Wait(dev_ctx_[p]);
+      }
+
       auto &var_name = static_cast<VarHandle *>(this->inputs_[0])->name_;
       VLOG(3) << "Invoke NCCL AllReduce";
       int dtype = -1;
-- 
GitLab