!3536 fix bug ArgMaxWithValue

Merge pull request !3536 from fangzehua/argmax
5 years ago · 38620e542a
parent e9316b1389 32519ef570
commit 38620e542a
1 changed files with 2 additions and 1 deletions
--- a/mindspore/ops/operations/array_ops.py
+++ b/mindspore/ops/operations/array_ops.py
@ -1185,7 +1185,8 @@ class ArgMaxWithValue(PrimitiveWithInfer):
        """init ArgMaxWithValue"""
        self.axis = axis
        self.keep_dims = keep_dims
-        _check_infer_attr_reduce(axis, keep_dims, self.name)
+        validator.check_value_type('keep_dims', keep_dims, [bool], self.name)
+        validator.check_value_type('axis', axis, [int], self.name)

    def infer_shape(self, x_shape):
        axis = self.axis