Add need_check_supported for maxpoolwithargmax.

4 years ago · 17b706c82d
parent f3453ef22c
commit 17b706c82d
2 changed files with 2 additions and 0 deletions
--- a/mindspore/ops/_op_impl/tbe/max_pool_grad_with_argmax.py
+++ b/mindspore/ops/_op_impl/tbe/max_pool_grad_with_argmax.py
@ -30,6 +30,7 @@ max_pool_grad_with_argmax_op_info = TBERegOp("MaxPoolGradWithArgmax") \
    .input(1, "grad", False, "required", "all") \
    .input(2, "argmax", False, "optional", "all") \
    .output(0, "y", False, "required", "all") \
+    .need_check_supported(True) \
    .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.U16_5HD, DataType.F16_5HD) \
    .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.I64_5HD, DataType.F16_5HD) \
    .get_op_info()
--- a/mindspore/ops/_op_impl/tbe/max_pool_with_argmax.py
+++ b/mindspore/ops/_op_impl/tbe/max_pool_with_argmax.py
@ -29,6 +29,7 @@ max_pool_with_argmax_op_info = TBERegOp("MaxPoolWithArgmax") \
    .input(0, "x", False, "required", "all") \
    .output(0, "y", False, "required", "all") \
    .output(1, "argmax", False, "required", "all") \
+    .need_check_supported(True) \
    .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.U16_5HD) \
    .get_op_info()