Inconsistent ONNX Export with Differentiable Head

Hello,

I am trying to fine-tune SetFit for a multi-class classification problem.

Everything is smooth until exporting to Onnx. The head is not exported correctly, so when loading the model with Onnx, the predictions are the outputs of some previous layer.

For example, I showcased this on a dummy dataset with 3 classes.

See https://colab.research.google.com/drive/19EESqbIDwD5FOI2Ufx22txFZ8qADq2xj?authuser=1#scrollTo=vl6AvQKqtU-9

This differs from the behaviour in the example notebook, where the LogisticRegression head is used.

Any directions would be appreciated; Otherwise, I would happily contribute with an MR if anyone can spot the issue.