I trained a HuggingFace Transformers XLNetForSequenceClassification model on a custom dataset with PyTorch backend. Then I used the provided convert_graph_to_onnx.py script to convert the model (from saved checkpoint) to ONNX format. I was able to load the model with ONNXRuntime, however when testing the model I get the following error.
When feeding in int64 numpy arrays input_ids and attention_masks, the model returns the following error except when both inputs have shape (x, 1) or (x, 6). There is nothing in the configuration of the model or the structure of the training data from my end that would require shape (x, 1) or (x, 6).
RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Add node. Name:'Add_26' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/math/element_wise_ops.h:361 void onnxruntime::BroadcastIterator::Init(int64_t, int64_t) axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 5 by 6

Any guidance on how to resolve this error would be much appreciated. Thank you!