Fix mode attribute value in DepthToSpace ONNX operation extractor (#3466)

This commit is contained in:
Anton Chetverikov 2020-12-03 19:33:07 +03:00 committed by GitHub
parent 58aee75914
commit a7ede592c3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -34,9 +34,9 @@ class DepthToSpaceFrontExtractor(FrontExtractorOp):
onnx_mode = onnx_attr(node, 'mode', 's', default=b'DCR').decode()
assert onnx_mode in ['DCR', 'CRD'], 'Unrecognized mode provided for DepthToSpace node {}'.format(node_name)
if onnx_mode == 'DCR':
mode = 'depth_first'
else:
mode = 'blocks_first'
else:
mode = 'depth_first'
DepthToSpaceOp.update_node_stat(node, {'block_size': block_size, 'mode': mode})
return cls.enabled