Fix whisper export (#1503)

* fix whisper config * make style --------- Co-authored-by: Mohit Sharma <[email protected]>
huggingface · Nov 2, 2023 · 1aee8ff · 1aee8ff
1 parent 87fcf9f
commit 1aee8ff
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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