diff --git a/optimum/exporters/onnx/model_patcher.py b/optimum/exporters/onnx/model_patcher.py index c534664010..db52493b6f 100644 --- a/optimum/exporters/onnx/model_patcher.py +++ b/optimum/exporters/onnx/model_patcher.py @@ -762,6 +762,8 @@ def patched_forward(input_ids, attention_mask): del result["input_ids"] if "attention_mask" in result: del result["attention_mask"] + if "all_layer_embeddings" in result: + del result["all_layer_embeddings"] return result