fix whisper config

huggingface · Nov 1, 2023 · 820e207 · 820e207
1 parent 9a9a7bf
commit 820e207
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