helinwang-patch-1
Yu Yang 7 years ago
parent f8141d90c8
commit 09935ab936

@ -345,8 +345,9 @@ struct NCCLAllReduceOpHandle : public OpHandle {
}
void Wait(platform::DeviceContext *waited_dev) override {
VLOG(3) << "Wait NCCL AllReduce";
this->dev_ctx_.at(waited_dev->GetPlace())->Wait();
for (auto &pair : member_->communication_streams_) {
pair.second.ctx_->Wait();
}
}
};

Loading…
Cancel
Save