Exporting to onnx model problem

[steven8274]

When I tried to export the 'VRT' model to onnx, the process was stopped by this error:

Traceback (most recent call last):
File "main_test_vrt.py", line 424, in
main()
File "main_test_vrt.py", line 71, in main
torch.onnx.export(model, # model being run
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/init.py", line 350, in export
return utils.export(
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 163, in export
_export(
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 1074, in _export
graph, params_dict, torch_out = _model_to_graph(
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 731, in _model_to_graph
graph = _optimize_graph(
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 308, in _optimize_graph
graph = _C._jit_pass_onnx(graph, operator_export_type)
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/init.py", line 416, in _run_symbolic_function
return utils._run_symbolic_function(*args, **kwargs)
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/utils.py", line 1406, in _run_symbolic_function
return symbolic_fn(g, *inputs, **attrs)
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 232, in wrapper
return fn(g, *args, **kwargs)
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_opset11.py", line 244, in pixel_shuffle
return symbolic_helper._unimplemented("pixel_shuffle", "only support 4d input")
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 440, in _unimplemented
_onnx_unsupported(f"{op}, {msg}")
File "/home/lxp/anaconda3/envs/hh_vrt/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 444, in _onnx_unsupported
raise RuntimeError(
RuntimeError: Unsupported: ONNX export of operator pixel_shuffle, only support 4d input. Please feel free to request support or submit a pull request on PyTorch GitHub.

From the code I know we use pytorch 1.8.1 or higher version to support '5D PixelShuffle', however, ONNX only support 4d input.So, how can I fix this problem?
Any help would be appreciated!