From a80afe8f1d4b48629db3312c981032f5c4dcf2a7 Mon Sep 17 00:00:00 2001 From: zhupengyang <1165938320@qq.com> Date: Thu, 13 Feb 2020 08:59:41 +0800 Subject: [PATCH] use default value if softmax has no 'axis' (#2857) --- lite/kernels/npu/bridges/softmax_op.cc | 2 +- lite/kernels/xpu/bridges/softmax_op.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/lite/kernels/npu/bridges/softmax_op.cc b/lite/kernels/npu/bridges/softmax_op.cc index 24bbb790e0..0ca3bc131d 100644 --- a/lite/kernels/npu/bridges/softmax_op.cc +++ b/lite/kernels/npu/bridges/softmax_op.cc @@ -42,7 +42,7 @@ int SoftmaxConverter(void* ctx, OpLite* op, KernelBase* kernel) { auto out_type = kernel->GetOutputDeclType("Out"); CHECK(out_type->precision() == PRECISION(kFloat)); CHECK(out_type->layout() == DATALAYOUT(kNCHW)); - auto axis = op_info->GetAttr("axis"); + int axis = op_info->HasAttr("axis") ? op_info->GetAttr("axis") : -1; if (axis < 0) { axis += x_rank; } diff --git a/lite/kernels/xpu/bridges/softmax_op.cc b/lite/kernels/xpu/bridges/softmax_op.cc index d964f29a86..7407640150 100644 --- a/lite/kernels/xpu/bridges/softmax_op.cc +++ b/lite/kernels/xpu/bridges/softmax_op.cc @@ -41,7 +41,7 @@ int SoftmaxConverter(void* ctx, OpLite* op, KernelBase* kernel) { auto out_type = kernel->GetOutputDeclType("Out"); CHECK(out_type->precision() == PRECISION(kFloat)); CHECK(out_type->layout() == DATALAYOUT(kNCHW)); - auto axis = op_info->GetAttr("axis"); + int axis = op_info->HasAttr("axis") ? op_info->GetAttr("axis") : -1; // X node std::shared_ptr x_node = nullptr; -- GitLab