diff --git a/lite/kernels/x86/reduce_op_function.h b/lite/kernels/x86/reduce_op_function.h index 05785bf009791f07871b68f9dbfed5ff4a7c9984..83570db1be182b6a60c3a0efe4227a5b45b55d38 100644 --- a/lite/kernels/x86/reduce_op_function.h +++ b/lite/kernels/x86/reduce_op_function.h @@ -46,40 +46,40 @@ void ReduceFunctor(const lite::Tensor& input, lite::Tensor* output, const std::vector& dims, bool keep_dim) { - auto x = EigenTensor::From(input); - - auto reduce_dim = Eigen::array(); - auto x_rank = static_cast(x.dimensions().size()); - for (size_t i = 0; i < dims.size(); ++i) { - if (dims[i] < 0) { - reduce_dim[i] = x_rank + dims[i]; - } else { - reduce_dim[i] = dims[i]; + auto te = strstr(typeid(Functor).name(), "SumFunctor"); + if (D == 3 && R_D == 1 && te != NULL) { + const lite::DDim& input_dims = input.dims(); + const T* input_data = input.data(); + T* output_data = output->mutable_data(); + for (int i = 0; i < input_dims[0]; i++) { + for (int k = 0; k < input_dims[2]; k++) { + int out_d = i * input_dims[2] + k; + T output_temp = 0; + for (int j = 0; j < input_dims[1]; j++) { + int input_d = + i * input_dims[1] * input_dims[2] + j * input_dims[2] + k; + output_temp = output_temp + input_data[input_d]; + } + output_data[out_d] = output_temp; + } } - } - - Functor functor; - if (D == 1) { - auto out = EigenScalar::From(output); - functor(&x, &out, reduce_dim); } else { - auto te = strstr(typeid(Functor).name(), "SumFunctor"); - if (D == 3 && R_D == 1 && te != NULL) { - const lite::DDim& input_dims = input.dims(); - const T* input_data = input.data(); - T* output_data = output->mutable_data(); - for (int i = 0; i < input_dims[0]; i++) { - for (int k = 0; k < input_dims[2]; k++) { - int out_d = i * input_dims[2] + k; - T output_temp = 0; - for (int j = 0; j < input_dims[1]; j++) { - int input_d = - i * input_dims[1] * input_dims[2] + j * input_dims[2] + k; - output_temp = output_temp + input_data[input_d]; - } - output_data[out_d] = output_temp; - } + auto x = EigenTensor::From(input); + + auto reduce_dim = Eigen::array(); + auto x_rank = static_cast(x.dimensions().size()); + for (size_t i = 0; i < dims.size(); ++i) { + if (dims[i] < 0) { + reduce_dim[i] = x_rank + dims[i]; + } else { + reduce_dim[i] = dims[i]; } + } + + Functor functor; + if (D == 1) { + auto out = EigenScalar::From(output); + functor(&x, &out, reduce_dim); } else { auto out = EigenTensor::From(*output, output->dims()); functor(&x, &out, reduce_dim);