mirrors
/
mmdeploy
mirror of https://github.com/open-mmlab/mmdeploy.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

compatible with torch 1.11 (#272)

pull/278/head
AllentDan 2022-03-25 23:21:10 +08:00 committed by GitHub
parent 3fa2cc8845
commit e01056cc0a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mmdeploy/apis/pytorch2onnx.py
View File

@ -32,6 +32,8 @@ def torch2onnx_impl(model: torch.nn.Module, input: torch.Tensor,
output_names = onnx_cfg['output_names']
axis_names = input_names + output_names
dynamic_axes = get_dynamic_axes(deploy_cfg, axis_names)
verbose = not onnx_cfg.get('strip_doc_string', True) or onnx_cfg.get(
'verbose', False)
# patch model
patched_model = patch_model(model, cfg=deploy_cfg, backend=backend)
@ -50,7 +52,7 @@ def torch2onnx_impl(model: torch.nn.Module, input: torch.Tensor,
dynamic_axes=dynamic_axes,
keep_initializers_as_inputs=onnx_cfg[
'keep_initializers_as_inputs'],
strip_doc_string=onnx_cfg.get('strip_doc_string', True))
verbose=verbose)
def torch2onnx(img: Any,