Wait all inputs ready

7 years ago · 631aa3d10a
parent 9b1f4d5d62
commit 631aa3d10a
1 changed files with 6 additions and 0 deletions
--- 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;