From dad7bdabd42ac2eeef7b3cb004ca64b6ad388cde Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 19 Mar 2018 20:17:32 +0800
Subject: [PATCH] Add setDev

---
 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 df04cfc46..c3a90149a 100644
--- a/paddle/fluid/framework/parallel_executor.cc
+++ b/paddle/fluid/framework/parallel_executor.cc
@@ -149,6 +149,7 @@ struct ScaleLossGradOpHandle : public OpHandle {
       auto stream =
           static_cast<platform::CUDADeviceContext *>(this->dev_ctx_[place_])
               ->stream();
+      cudaSetDevice(boost::get<platform::CUDAPlace>(place_).device);
       VLOG(3) << "1";
       PADDLE_ENFORCE(cudaGetLastError());
       VLOG(3) << "2";
@@ -163,7 +164,7 @@ struct ScaleLossGradOpHandle : public OpHandle {
 
   void Wait(platform::DeviceContext *waited_dev) override {
     if (platform::is_cpu_place(waited_dev->GetPlace())) {
-      this->dev_ctx_.at(place_)->Wait();
+      dev_ctx_.at(place_)->Wait();
     } else {
       auto stream =
           static_cast<platform::CUDADeviceContext *>(waited_dev)->stream();
-- 
GitLab