onnx output shape is wrong

[CallmeZhangChenchen]

System Info

optimum  1.24.0

Who can help?

@michaelbenayoun

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer

model_checkpoint = "moka-ai/m3e-base"
save_directory = "test"

# Load a model from transformers and export it to ONNX
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
import pdb;pdb.set_trace()
ort_model = ORTModelForFeatureExtraction.from_pretrained(model_checkpoint, export=True, from_transformers=True)

# Save the ONNX model and tokenizer
ort_model.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)

onnx output
{'name': 'last_hidden_state',
'shape': ['batch_size', 'sequence_length', 768],
'type': 'tensor(float)'}

Expected behavior

onnx output
{'name': 'last_hidden_state',
'shape': ['batch_size', 1, 768],
'type': 'tensor(float)'}