diff --git a/paddle/fluid/operators/sum_op.h b/paddle/fluid/operators/sum_op.h index 34403c7a7aa717cca470be2931009e219e00e3ae..11987c61aebaad00f8a71f1b909c83c44ddc8b0e 100644 --- a/paddle/fluid/operators/sum_op.h +++ b/paddle/fluid/operators/sum_op.h @@ -43,17 +43,31 @@ class SumKernel : public framework::OpKernel { out->mutable_data(context.GetPlace()); } auto result = EigenVector::Flatten(*out); + auto &place = + *context.template device_context().eigen_device(); + int start = in_place ? 1 : 0; if (!in_place) { - math::SetConstant constant_functor; - constant_functor(context.template device_context(), out, - 0.0); + if ((in_num >= 2) && in_vars[0]->IsType() && + in_vars[1]->IsType()) { + auto &in_0 = in_vars[0]->Get(); + auto &in_1 = in_vars[1]->Get(); + if (in_0.numel() && in_1.numel()) { + auto in_0_e = EigenVector::Flatten(in_0); + auto in_1_e = EigenVector::Flatten(in_1); + result.device(place) = in_0_e + in_1_e; + start = 2; + } + } + if (start != 2) { + math::SetConstant constant_functor; + constant_functor(context.template device_context(), + out, 0.0); + } } math::SelectedRowsAddToTensor functor; - auto &place = - *context.template device_context().eigen_device(); // If in_place, just skip the first tensor - for (size_t i = in_place ? 1 : 0; i < in_num; i++) { + for (size_t i = start; i < in_num; i++) { if (in_vars[i]->IsType()) { auto &in_t = in_vars[i]->Get(); if (in_t.numel() == 0) {