From 1aee8ff05c9aca4877c28109e8eadf5a47d5539c Mon Sep 17 00:00:00 2001
From: Mohit Sharma <mohit21sharma.ms@gmail.com>
Date: Thu, 2 Nov 2023 13:43:26 +0530
Subject: [PATCH] Fix whisper export (#1503)

* fix whisper config

* make style

---------

Co-authored-by: Mohit Sharma <mohit@huggingface.co>
---
 optimum/exporters/onnx/model_configs.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
index 5b1268f3c64..2da3f5bea6b 100644
--- a/optimum/exporters/onnx/model_configs.py
+++ b/optimum/exporters/onnx/model_configs.py
@@ -1291,7 +1291,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]:
 
 
 class WhisperOnnxConfig(AudioToTextOnnxConfig):
-    NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig
+    NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args(
+        encoder_num_layers="encoder_layers",
+        decoder_num_layers="decoder_layers",
+    )
     ATOL_FOR_VALIDATION = 1e-3
 
     @property