提交 38620e54 编写于 作者: M mindspore-ci-bot 提交者: Gitee

!3536 fix bug ArgMaxWithValue

Merge pull request !3536 from fangzehua/argmax
......@@ -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
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册