Repository: huggingface/transformers (Python)
Task type: Bug Fix

Issue:
Not handled case when use_weighted_layer_sum and return-dict=True in WhisperForAudioClassification
@sanchit-gandhi 
I use WhisperForAudioClassification task and want to use `use_weighted_layer_sum=True`, but there is a problem when call forward, 
the encoder part can return tuple or dict if `return_dict=True` but the code for use `use_weighted_layer_sum=True` assume the return to be tuple only and this line raise error `hidden_states = torch.stack(encoder_outputs, dim=1)` if the encoder return dict, there are workaround by using `return_dict=False` but when use the model later with `pipeline` it will raise error because it assume the model to return dict not tuple. [Link to code with the problem](https://github.com/huggingface/transformers/blob/c7f076a00ee54f777b3d3322c91bc11489a47950/src/transformers/models/whisper/modeling_whisper.py#L2918C6-L2918C6)
```py
        if self.config.use_weighted_layer_sum:
            hidden_states = torch.stack(encoder_outputs, dim=1) # This line raise error when return_dict=True and use_weighted_layer_sum=True
            norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
            hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
        else:
            hidden_states = encoder_outputs[0]
```
**Reproduce error**
```py
from transformers import  WhisperForAudioClassification, AutoFeatureExtractor
from datasets import load_dataset

dataset = load_dataset('seba3y/speechocean762',)
dataset = dataset['train']
sampling_rate = dataset.features["audio"].sampling_rate
dataset = dataset.remove_columns(['utt_name', 'text', 'completeness', 'fluency', 'prosodic'])


feature_extractor = AutoFeatureExtractor.from_pretrained("seba3y/whisper-tiny")
model = WhisperForAudioClassification.from_pretrained("seba3y/whisper-tiny",
                                                       use_weighted_layer_sum=True, 
                                                       return_dict=True)
# test if it work
inputs = feature_extractor(dataset['train'][3]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
with torch.no_grad():
    logits = model(**inputs).logits
predicted_class_ids = torch.argmax(logits, dim=-1).item()
predicted_label = model.config.id2label[predicted_class_ids]
print(predicted_label)
```












Discussion:
Hi @ElsebaiyMohamed, thanks for raising this issue and providing details on the error + a snippet. Could you also provide information about the running environment: run `transformers-cli env` in the terminal and copy-paste the output? 
Hi @amyeroberts ,
Apologies for the delayed response! 🙏 Life threw a curveball, but I'm back on track. Thanks for your patience!

Regarding your request, here's the output of `transformers-cli env`:

```bash
transformers version: 4.36.0
Platform: Linux-5.15.133+-x86_64-with-glibc2.35
Python version: 3.10.12
Huggingface_hub version: 0.19.4
Safetensors version: 0.4.1
Accelerate version: 0.25.0
Accelerate config: 	not found
PyTorch version (GPU?): 2.0.0 (True)
Tensorflow version (GPU?): 2.13.0 (True)
Flax version (CPU?/GPU?/TPU?): 0.7.5 (gpu)
Jax version: 0.4.21
JaxLib version: 0.4.21
Using GPU in script?: yes
Using distributed or parallel set-up in script?: no
```

Let me know if there's anything else I can help you with.
@ElsebaiyMohamed Great - thanks for providing this info! 

cc @sanchit-gandhi @ylacombe 

Solution patch:
diff --git a/src/transformers/models/whisper/modeling_whisper.py b/src/transformers/models/whisper/modeling_whisper.py
--- a/src/transformers/models/whisper/modeling_whisper.py
+++ b/src/transformers/models/whisper/modeling_whisper.py
@@ -57,6 +57,8 @@
 
 logger = logging.get_logger(__name__)
 
+_HIDDEN_STATES_START_POSITION = 1
+
 _CONFIG_FOR_DOC = "WhisperConfig"
 _CHECKPOINT_FOR_DOC = "openai/whisper-tiny"
 
@@ -2957,6 +2959,11 @@ def forward(
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
+        if self.config.use_weighted_layer_sum:
+            output_hidden_states = True
+        elif output_hidden_states is None:
+            output_hidden_states = self.config.output_hidden_states
+
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if encoder_outputs is None:
@@ -2969,7 +2976,8 @@ def forward(
             )
 
         if self.config.use_weighted_layer_sum:
-            hidden_states = torch.stack(encoder_outputs, dim=1)
+            hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION]
+            hidden_states = torch.stack(hidden_states, dim=1)
             norm_weights = nn.functional.softmax(self.layer_weights, dim=-1)
             hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
         else:


Test patch:
diff --git a/tests/models/whisper/test_modeling_whisper.py b/tests/models/whisper/test_modeling_whisper.py
--- a/tests/models/whisper/test_modeling_whisper.py
+++ b/tests/models/whisper/test_modeling_whisper.py
@@ -2292,16 +2292,15 @@ def get_subsampled_output_lengths(self, input_lengths):
     def encoder_seq_length(self):
         return self.get_subsampled_output_lengths(self.seq_length)
 
-    def create_and_check_model_forward(self, config, inputs_dict, freeze_encoder=False):
-        model = WhisperForAudioClassification(config=config).to(torch_device).eval()
-
-        if freeze_encoder:
-            model.freeze_encoder()
+    def create_and_check_model_forward(self, config, inputs_dict, use_weighted_layer_sum=False):
+        config.use_weighted_layer_sum = use_weighted_layer_sum
+        model = WhisperForAudioClassification(config=config)
+        model.to(torch_device).eval()
 
         input_features = inputs_dict["input_features"]
 
-        # first forward pass
-        last_hidden_state = model(input_features).logits
+        with torch.no_grad():
+            last_hidden_state = model(input_features).logits
 
         self.parent.assertTrue(last_hidden_state.shape, (13, 2))
 
@@ -2336,6 +2335,14 @@ def test_forward_signature(self):
             expected_arg_names = ["input_features", "head_mask", "encoder_outputs"]
             self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names)
 
+    def test_forward_pass(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_model_forward(*config_and_inputs)
+
+    def test_forward_pass_weighted_layer_sum(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_model_forward(*config_and_inputs, use_weighted_layer_sum=True)
+
     @unittest.skip(reason="Some undefined behavior encountered with tiny versions of this model. Skip for now.")
     def test_cpu_offload(self):
         pass