fix bug of reduce ops infer value

2021-06-23 15:14:48 +08:00 · 2021-06-23 15:14:48 +08:00 · c336870af2
parent f98497ca09
commit c336870af2
1 changed files with 1 additions and 1 deletions
--- a/mindspore/ops/operations/math_ops.py
+++ b/mindspore/ops/operations/math_ops.py
@ -352,7 +352,7 @@ class _Reduce(PrimitiveWithInfer):

            if np_reduce_func is not None:
                value = input_x['value'].asnumpy()
-                if not axis_v:
+                if not axis_v and axis_v != 0:
                    axis_v = [i for i in range(len(input_x['shape']))]
                    axis_v = tuple(axis_v)
                value = np_reduce_func(value, axis_v, keepdims=self.keep_dims)