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