diff options
Diffstat (limited to 'model-optimizer/extensions/front/onnx/flatten_ext.py')
| -rw-r--r-- | model-optimizer/extensions/front/onnx/flatten_ext.py | 8 |
1 files changed, 3 insertions, 5 deletions
diff --git a/model-optimizer/extensions/front/onnx/flatten_ext.py b/model-optimizer/extensions/front/onnx/flatten_ext.py index 1296e5b66..11aaa1b9d 100644 --- a/model-optimizer/extensions/front/onnx/flatten_ext.py +++ b/model-optimizer/extensions/front/onnx/flatten_ext.py @@ -15,9 +15,8 @@ """ from mo.front.extractor import FrontExtractorOp -from mo.ops.op import Op - from mo.front.onnx.extractors.utils import onnx_attr +from mo.ops.flatten_onnx import FlattenONNX class FlattenFrontExtractor(FrontExtractorOp): @@ -31,6 +30,5 @@ class FlattenFrontExtractor(FrontExtractorOp): 'axis': axis } - # update the attributes of the node - Op.get_op_class_by_name(__class__.op).update_node_stat(node, attrs) - return __class__.enabled
\ No newline at end of file + FlattenONNX.update_node_stat(node, attrs) + return __class__.enabled |