From bade579826d0e6e82b62b6f0b630dbfee35f65d2 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 19 Mar 2018 20:08:52 +0800
Subject: [PATCH] Wait code

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

diff --git a/paddle/fluid/framework/parallel_executor.cc b/paddle/fluid/framework/parallel_executor.cc
index 94c61461c0..bc9035b302 100644
--- a/paddle/fluid/framework/parallel_executor.cc
+++ b/paddle/fluid/framework/parallel_executor.cc
@@ -193,7 +193,8 @@ struct FetchOpHandle : public OpHandle {
 
   void Run() override {
     for (auto *input : inputs_) {
-      input->generated_op_->Wait(nullptr);
+      auto *var = static_cast<VarHandle *>(input);
+      var->generated_op_->Wait(this->dev_ctx_[var->place_]);
     }
 
     tensors_.resize(inputs_.size());
-- 
GitLab