From a326e62ed400172c7a400cb84c44596e8a0d30a8 Mon Sep 17 00:00:00 2001 From: fangzehua Date: Tue, 12 Apr 2022 12:55:49 +0000 Subject: [PATCH] fix stack op infer shape --- mindspore/core/abstract/prim_arrays.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindspore/core/abstract/prim_arrays.cc b/mindspore/core/abstract/prim_arrays.cc index ba65e6adc2c..58002ea6c53 100644 --- a/mindspore/core/abstract/prim_arrays.cc +++ b/mindspore/core/abstract/prim_arrays.cc @@ -149,7 +149,7 @@ AbstractBasePtr InferImplStack(const AnalysisEnginePtr &, const PrimitivePtr &pr ValuePtr axis = primitive->GetAttr("axis"); // Axis value should be in [-(rank_base + 1), rank_base). - int64_t axis_value = CheckAxis(op_name, "axis", axis, -(rank_base + 1), rank_base, "input_x"); + int64_t axis_value = CheckAxis(op_name, "axis", axis, -(rank_base + 1), rank_base + 1, "input_x"); for (size_t i = 1; i < tuple_len; ++i) { AbstractTensorPtr tensor = nullptr;