diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 2fb2a5c2952..6c7dec773c4 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -2124,9 +2124,9 @@ def outputs(self) -> Dict[str, Dict[int, str]]: # for Speech2text, we need to name the second axis as # encoder_sequence_length / 2 * self._config.num_conv_layers as the axis name is # used for dummy input generation - common_outputs["last_hidden_state"][1] = ( - f"{common_outputs['last_hidden_state'][1]} / {(2 * self._config.num_conv_layers)}" - ) + common_outputs["last_hidden_state"][ + 1 + ] = f"{common_outputs['last_hidden_state'][1]} / {(2 * self._config.num_conv_layers)}" return common_outputs