faceformer export to onnx failed

I want to export vocaset.pth to onnx model with the following rectify :

firstly , in demo.py added the export code:

input_names = ['audio_feature', 'template', 'one_hot']
output_names = ['vertice_out']
torch.onnx.export(model,                        # model being run
            (audio_feature, template, one_hot),                              # model input (or a tuple for multiple inputs)
            'vocaset.onnx',                # where to save the model (can be a file or file-like object)
            export_params=True,             # store the trained parameter weights inside the model file
            opset_version=11,               # the ONNX version to export the model to
            do_constant_folding=True,       # whether to execute constant folding for optimization
            input_names = input_names,        # the model's input names
            output_names = output_names,      # the model's output names
            dynamic_axes={'audio_feature' : {1 : 'audio_len'}}
            )

and secondly, rewrite the forward function same to the predict function in faceformer.py. when I running the demo.py, it failed with the following message:

(function ComputeConstantFolding) [W shape_type_inference.cpp:419] Warning: Constant folding in symbolic shape inference fails: expected scalar type Long but found Float Exception raised from data_ptr at /pytorch/build/aten/src/ATen/core/TensorMethods.cpp:5759 (most recent call first): frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x42 (0x7fd3b3eafa22 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libc10.so) frame #1: c10::detail::torchCheckFail(char const, char const, unsigned int, std::string const&) + 0x5b (0x7fd3b3eac3db in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libc10.so) frame #2: long at::Tensor::data_ptr() const + 0xde (0x7fd22608d83e in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so) frame #3: torch::jit::onnx_constant_fold::runTorchSlice_opset10(torch::jit::Node const, std::vector<at::Tensor, std::allocator >&) + 0x42e (0x7fd36aa778fe in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #4: torch::jit::onnx_constant_fold::runTorchBackendForOnnx(torch::jit::Node const, std::vector<at::Tensor, std::allocator >&, int) + 0x1c5 (0x7fd36aa78c45 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #5: + 0xafd7f1 (0x7fd36aab77f1 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #6: torch::jit::ONNXShapeTypeInference(torch::jit::Node, std::map<std::string, c10::IValue, std::less, std::allocator<std::pair<std::string const, c10::IValue> > > const&, int) + 0x906 (0x7fd36aabc666 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #7: + 0xb05414 (0x7fd36aabf414 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #8: + 0xa7c010 (0x7fd36aa36010 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #9: + 0x500f98 (0x7fd36a4baf98 in /opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/lib/libtorch_python.so) frame #10: _PyMethodDef_RawFastCallKeywords + 0x254 (0x564c0405c7a4 in /opt/conda/envs/faceformer/bin/python) frame #11: + 0x17fb40 (0x564c04092b40 in /opt/conda/envs/faceformer/bin/python) frame #12: _PyEval_EvalFrameDefault + 0x4762 (0x564c040da702 in /opt/conda/envs/faceformer/bin/python) frame #13: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #14: _PyFunction_FastCallKeywords + 0x583 (0x564c0404bcd3 in /opt/conda/envs/faceformer/bin/python) frame #15: + 0x17f9c5 (0x564c040929c5 in /opt/conda/envs/faceformer/bin/python) frame #16: _PyEval_EvalFrameDefault + 0x1401 (0x564c040d73a1 in /opt/conda/envs/faceformer/bin/python) frame #17: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #18: _PyFunction_FastCallKeywords + 0x583 (0x564c0404bcd3 in /opt/conda/envs/faceformer/bin/python) frame #19: + 0x17f9c5 (0x564c040929c5 in /opt/conda/envs/faceformer/bin/python) frame #20: _PyEval_EvalFrameDefault + 0x1401 (0x564c040d73a1 in /opt/conda/envs/faceformer/bin/python) frame #21: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #22: _PyFunction_FastCallKeywords + 0x583 (0x564c0404bcd3 in /opt/conda/envs/faceformer/bin/python) frame #23: + 0x17f9c5 (0x564c040929c5 in /opt/conda/envs/faceformer/bin/python) frame #24: _PyEval_EvalFrameDefault + 0x1401 (0x564c040d73a1 in /opt/conda/envs/faceformer/bin/python) frame #25: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #26: _PyFunction_FastCallKeywords + 0x521 (0x564c0404bc71 in /opt/conda/envs/faceformer/bin/python) frame #27: + 0x17f9c5 (0x564c040929c5 in /opt/conda/envs/faceformer/bin/python) frame #28: _PyEval_EvalFrameDefault + 0x4762 (0x564c040da702 in /opt/conda/envs/faceformer/bin/python) frame #29: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #30: _PyFunction_FastCallKeywords + 0x583 (0x564c0404bcd3 in /opt/conda/envs/faceformer/bin/python) frame #31: + 0x17f9c5 (0x564c040929c5 in /opt/conda/envs/faceformer/bin/python) frame #32: _PyEval_EvalFrameDefault + 0x1401 (0x564c040d73a1 in /opt/conda/envs/faceformer/bin/python) frame #33: _PyFunction_FastCallDict + 0x118 (0x564c0404acf8 in /opt/conda/envs/faceformer/bin/python) frame #34: _PyEval_EvalFrameDefault + 0x1cb8 (0x564c040d7c58 in /opt/conda/envs/faceformer/bin/python) frame #35: _PyEval_EvalCodeWithName + 0xdf9 (0x564c0402ca29 in /opt/conda/envs/faceformer/bin/python) frame #36: _PyFunction_FastCallKeywords + 0x583 (0x564c0404bcd3 in /opt/conda/envs/faceformer/bin/python) frame #37: _PyEval_EvalFrameDefault + 0x3f5 (0x564c040d6395 in /opt/conda/envs/faceformer/bin/python) frame #38: _PyFunction_FastCallKeywords + 0x187 (0x564c0404b8d7 in /opt/conda/envs/faceformer/bin/python) frame #39: _PyEval_EvalFrameDefault + 0x3f5 (0x564c040d6395 in /opt/conda/envs/faceformer/bin/python) frame #40: _PyEval_EvalCodeWithName + 0x255 (0x564c0402be85 in /opt/conda/envs/faceformer/bin/python) frame #41: PyEval_EvalCode + 0x23 (0x564c0402d273 in /opt/conda/envs/faceformer/bin/python) frame #42: + 0x227c82 (0x564c0413ac82 in /opt/conda/envs/faceformer/bin/python) frame #43: PyRun_FileExFlags + 0x9e (0x564c04144e1e in /opt/conda/envs/faceformer/bin/python) frame #44: PyRun_SimpleFileExFlags + 0x1bb (0x564c0414500b in /opt/conda/envs/faceformer/bin/python) frame #45: + 0x2330fa (0x564c041460fa in /opt/conda/envs/faceformer/bin/python) frame #46: _Py_UnixMain + 0x3c (0x564c0414618c in /opt/conda/envs/faceformer/bin/python) frame #47: __libc_start_main + 0xe7 (0x7fd3fae99c87 in /lib/x86_64-linux-gnu/libc.so.6) frame #48: + 0x1d803a (0x564c040eb03a in /opt/conda/envs/faceformer/bin/python) (function ComputeConstantFolding) Traceback (most recent call last): File "/workspace/FaceFormer/demo.py", line 250, in main() File "/workspace/FaceFormer/demo.py", line 246, in main test_model(args) File "/opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context return func(*args, **kwargs) File "/workspace/FaceFormer/demo.py", line 99, in test_model output_names = output_names, # the model's output names File "/opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/onnx/init.py", line 280, in export custom_opsets, enable_onnx_checker, use_external_data_format) File "/opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/onnx/utils.py", line 94, in export use_external_data_format=use_external_data_format) File "/opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/onnx/utils.py", line 695, in _export dynamic_axes=dynamic_axes) File "/opt/conda/envs/faceformer/lib/python3.7/site-packages/torch/onnx/utils.py", line 502, in _model_to_graph _export_onnx_opset_version) RuntimeError: expected scalar type Long but found Float

does anyone meet the same problem like me? need your help! tks!!

EvelynFan / FaceFormer

faceformer export to onnx failed #39