diff --git a/paddle/fluid/operators/elementwise_add_op.cc b/paddle/fluid/operators/elementwise_add_op.cc index d2c20537136fc3ac9d1bece24a2238f26215c922..c1ddc1824bb0531a7e628d858bd2d3b0e4718877 100644 --- a/paddle/fluid/operators/elementwise_add_op.cc +++ b/paddle/fluid/operators/elementwise_add_op.cc @@ -25,6 +25,6 @@ REGISTER_OP_CPU_KERNEL( REGISTER_OP_CPU_KERNEL( elementwise_add_grad, ops::ElementwiseAddGradKernel, - ops::ElementwiseAddGradKernel, - ops::ElementwiseAddGradKernel, - ops::ElementwiseAddGradKernel); + ops::ElementwiseAddGradKernel); +// ops::ElementwiseAddGradKernel, +// ops::ElementwiseAddGradKernel); diff --git a/paddle/fluid/operators/elementwise_add_op.h b/paddle/fluid/operators/elementwise_add_op.h index d75d86c242d3f06db51fd3623a9d42e944791a96..5984f4aef1d92504eb369020fb4503b1a04cd9f0 100644 --- a/paddle/fluid/operators/elementwise_add_op.h +++ b/paddle/fluid/operators/elementwise_add_op.h @@ -98,9 +98,22 @@ class ElementwiseAddGradKernel : public framework::OpKernel { auto* dx = ctx.Output(framework::GradVarName("X")); auto* dy = ctx.Output(framework::GradVarName("Y")); int axis = ctx.Attr("axis"); - ElemwiseGradCompute, IdentityGrad>( - ctx, *x, *y, *out, *dout, axis, dx, dy, IdentityGrad(), - IdentityGrad()); + + if (platform::is_cpu_place(ctx.GetPlace()) && (x->dims() == y->dims())) { + auto blas = math::GetBlas(ctx); + + if (dx) + dx->mutable_data(ctx.GetPlace()); + if (dy) + dy->mutable_data(ctx.GetPlace()); + + blas.VCOPY(dout->numel(), dout->data(), dx->data()); + blas.VCOPY(dout->numel(), dout->data(), dy->data()); + } else { + ElemwiseGradCompute, IdentityGrad>( + ctx, *x, *y, *out, *dout, axis, dx, dy, IdentityGrad(), + IdentityGrad()); + } } };