torch模型导出到onnx问题 expected 1 dims but tensor has 2

导出过程中出现如下错误

expected 1 dims but tensor has 2

torch 1.5.0 must be used for the export because the following error occurs intorch 1.6.0 caused by the internal optimization process.
File “/usr/local/lib/python3.7/site-packages/torch/onnx/utils.py”, line 409, in _model_to_graph
export_onnx_opset_version)
RuntimeError: expected 1 dims but tensor has 2
==============================================

参考出处
https://medium.com/axinc-ai/st-gcn-a-machine-learning-model-for-detecting-human-actions-from-skeletons-46a95b31b5db

核心就是版本问题，换个版本就解决了。

推理和训练阶段可以写的不一样，推理尽可能的整合为好。