Browse Source

fix argmaxwithvalue

tags/v0.7.0-beta
fangzehua 5 years ago
parent
commit
32519ef570
1 changed files with 2 additions and 1 deletions
  1. +2
    -1
      mindspore/ops/operations/array_ops.py

+ 2
- 1
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


Loading…
Cancel
Save