Issue with converting Comet model to ONNX. Split-node error.
Labels: module: onnx, triaged

I'm trying to convert the [Comet](https://unbabel.github.io/COMET/html/index.html) model to ONNX. In particular I'm working with the referenceless version of the metric ("wmt21-comet-qe-mqm"). It is based on BERT and contains a regression head.

The export to ONNX succeeds. And the resulting model passes the checks. However, performing inference fails with an error - Cannot split using values in 'split' attribute. Below I provide details of the error.

I wonder if the ONNX supports such configuration. I also prepared a short script that may be used to reproduce this issue.

```python
import os
from transformers import AutoTokenizer
import torch
from comet import download_model, load_from_checkpoint
import onnxruntime as ort
import onnx

model = load_from_checkpoint(download_model("wmt21-comet-qe-mqm"))

CONVERTED_COMET_DIR = 'CONVERTED_COMET'
CONVERTED_COMET_MODEL_PATH = os.path.join(CONVERTED_COMET_DIR, 'model.onnx')
tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-large')
MODEL_MAX_LENGTH = tokenizer.model_max_length

try: 
    os.makedirs(CONVERTED_COMET_DIR)
except FileExistsError:
    pass

input_names = ["src_input_ids", "src_attention_mask", "mt_input_ids", "mt_attention_mask"]
inputs = {key: torch.ones(1, MODEL_MAX_LENGTH, dtype=torch.int64) for key in input_names}
symbolic_names = {0: "batch_size"}

torch.onnx.export(
    model,
    (*[inputs[key] for key in input_names],),
    CONVERTED_COMET_MODEL_PATH,
    opset_version=15,
    do_constant_folding=True,
    input_names=input_names,
    output_names=["score"],
    dynamic_axes={key: symbolic_names for key in input_names},
)

onnx.checker.check_model(CONVERTED_COMET_MODEL_PATH)
ort_sess = ort.InferenceSession(CONVERTED_COMET_MODEL_PATH)

input_src = tokenizer("Algo está mal aquí...", return_tensors="np", max_length=MODEL_MAX_LENGTH, padding='max_length')
input_mt = tokenizer("Coś tu nie gra...", return_tensors="np", max_length=MODEL_MAX_LENGTH, padding='max_length')
inp = {"src_input_ids": input_src['input_ids'],
       "src_attention_mask": input_src['attention_mask'],
       "mt_input_ids": input_mt['input_ids'],
       "mt_attention_mask": input_mt['attention_mask']}

outputs = ort_sess.run(None, inp)
```

The error details
```
2022-09-22 12:34:13.957212579 [E:onnxruntime:, sequential_executor.cc:368 Execute] Non-zero status code returned while running Split node. Name:'Split_5573' Status Message: Cannot split using values in 'split' attribute. Axis=0 Input shape={1,512} NumOutputs=1 Num entries in 'split' (must equal number of outputs) was 1 Sum of sizes in 'split' (must equal size of selected axis) was 8
---------------------------------------------------------------------------
Fail                                      Traceback (most recent call last)
/tmp/ipykernel_7700/1124372680.py in <module>
     60 inp = prepare_pair("Algo está mal aquí...", "Coś tu nie gra...")
     61 
---> 62 outputs = ort_sess.run(None, inp)

/opt/conda/lib/python3.7/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py in run(self, output_names, input_feed, run_options)
    198             output_names = [output.name for output in self._outputs_meta]
    199         try:
--> 200             return self._sess.run(output_names, input_feed, run_options)
    201         except C.EPFail as err:
    202             if self._enable_fallback:

Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Split node. Name:'Split_5573' Status Message: Cannot split using values in 'split' attribute. Axis=0 Input shape={1,512} NumOutputs=1 Num entries in 'split' (must equal number of outputs) was 1 Sum of sizes in 'split' (must equal size of selected axis) was 8
```

### Versions

Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (Debian 8.3.0-6) 8.3.0
Clang version: Could not collect
CMake version: version 3.13.4
Libc version: glibc-2.10

Python version: 3.7.12 | packaged by conda-forge | (default, Oct 26 2021, 06:08:53)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-4.19.0-21-cloud-amd64-x86_64-with-debian-10.12
Is CUDA available: False
CUDA runtime version: 11.3.109
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.6
[pip3] pytorch-lightning==1.6.4
[pip3] torch==1.11.0
[pip3] torch-xla==1.11
[pip3] torchmetrics==0.8.2
[pip3] torchvision==0.12.0+cu113
[conda] blas                      2.115                       mkl    conda-forge
[conda] blas-devel                3.9.0            15_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] dlenv-pytorch-1-11-gpu    1.0.20220630     py37hc1c1d6d_0    file:///tmp/conda-pkgs
[conda] libblas                   3.9.0            15_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            15_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            15_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            15_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.21.6                   pypi_0    pypi
[conda] pytorch                   1.11.0          py3.7_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-lightning         1.6.4                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchmetrics              0.8.2                    pypi_0    pypi
[conda] torchvision               0.12.0+cu113             pypi_0    pypi

Additional libraries:
unbabel-comet==1.1.2
transformers==4.17.0