remove unwanted output

huggingface · Jan 9, 2024 · f782fda · f782fda
1 parent 7f92241
commit f782fda
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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