Add remove_value_output attribute to onnx ArgMax extractor (#4831)

* add remove_value_output to onnx argmax extractor * update TopK extender * update copyright * change condition in topk extender based on connection logic
2021-03-22 14:19:27 +03:00
parent 1f4f449308
commit 244af264e6
2 changed files with 5 additions and 3 deletions
--- a/model-optimizer/extensions/front/onnx/argmax_ext.py
+++ b/model-optimizer/extensions/front/onnx/argmax_ext.py
@@ -35,7 +35,9 @@ class ArgMaxFrontExtractor(FrontExtractorOp):
            'out_max_val' : 0,

            # Set attribute to trigger ArgMax replacer in case do not keep the dimension
-            'keepdims': keepdims
+            'keepdims': keepdims,
+
+            'remove_values_output': True
        }

        ArgMaxOp.update_node_stat(node, attrs)