pj3fizlut
/
mindspore
forked from OSSInnovation/mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

!3536 fix bug ArgMaxWithValue 

Merge pull request !3536 from fangzehua/argmax
This commit is contained in:
mindspore-ci-bot 2020-07-27 19:33:09 +08:00 committed by Gitee
parent e9316b1389 32519ef570
commit 38620e542a
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mindspore/ops/operations/array_ops.py
View File

@ -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