From 29cc9f308d151c23ddbaeef69530f3c7c56a6ce4 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 19 Mar 2018 19:39:13 +0800
Subject: [PATCH] SetDev for nccl

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

diff --git a/paddle/fluid/framework/parallel_executor.cc b/paddle/fluid/framework/parallel_executor.cc
index 6777aec488d..f7dc8339371 100644
--- a/paddle/fluid/framework/parallel_executor.cc
+++ b/paddle/fluid/framework/parallel_executor.cc
@@ -358,7 +358,7 @@ struct NCCLAllReduceOpHandle : public OpHandle {
         }
 
         auto &nccl_ctx = member_->communication_streams_.at(dev_id);
-
+        cudaSetDevice(dev_id);
         platform::dynload::ncclAllReduce(
             buffer, buffer, numel, static_cast<ncclDataType_t>(dtype), ncclSum,
             nccl_ctx.comm, nccl_ctx.stream());
@@ -519,7 +519,6 @@ void ParallelExecutor::ConstructDependencyGraph(
             var.name_ = og;
             var.version_ = vars.size() - 1;
             op_handle->outputs_.emplace_back(&var);
-
             op_handle->dev_ctx_[p] = member_->CommunicationDevCtx(p);
           }
         }
-- 
GitLab