Mismatched dimensions error when using beam_size != 1 and return_timestamps="word" for AutomaticSpeechRecognitionPipeline with Whisper
Author: nestormh
Pull request: no

### System Info

- `transformers` version: 4.33.3
- Platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.31
- Python version: 3.9.18
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No


### Who can help?

@sanchit-gandhi

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I have the following code:
```
model_name_or_path = "openai/whisper-small"
LANGUAGE = "French"
TASK = "transcribe"
CHUNK_LENGTH_S = 30  # 30 seconds chunks
SAMPLE_RATE = 16_000

model = WhisperForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path=model_name_or_path,
    load_in_8bit=True,
    device_map="cuda:0",
)

feature_extractor = WhisperFeatureExtractor.from_pretrained(
    model_name_or_path,
)
tokenizer = WhisperTokenizer.from_pretrained(
    model_name_or_path,
    language=LANGUAGE,
    task=TASK,
)
processor = WhisperProcessor.from_pretrained(
    model_name_or_path,
    language=LANGUAGE,
    task=TASK,
)

pipeline_generator = AutomaticSpeechRecognitionPipeline(
    model=model,
    tokenizer=tokenizer,
    feature_extractor=feature_extractor,
    chunk_length_s=CHUNK_LENGTH_S,
    stride_length_s=5
)
pipeline_generator.model.config.forced_decoder_ids = (
    pipeline_generator.tokenizer.get_decoder_prompt_ids(
        language=LANGUAGE, task=TASK
    )
)

forced_decoder_ids = processor.get_decoder_prompt_ids(
    language=LANGUAGE, task=TASK
)

model.config.use_cache = True
model.eval()

transcript = pipeline_generator(
    inputs=inputs,
    return_timestamps="word",
    generate_kwargs={
        "num_beams": 5,
    },
)
```

After executing the code, I receive the following error:
```
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/automatic_speech_recognition.py", line 356, in __call__
    return super().__call__(inputs, **kwargs)
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/base.py", line 1132, in __call__
    return next(
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 124, in __next__
    item = next(self.iterator)
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/pt_utils.py", line 266, in __next__
    processed = self.infer(next(self.iterator), **self.params)
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/base.py", line 1046, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/pipelines/automatic_speech_recognition.py", line 551, in _forward
    tokens = self.model.generate(
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/models/whisper/modeling_whisper.py", line 1768, in generate
    outputs["token_timestamps"] = self._extract_token_timestamps(outputs, generation_config.alignment_heads)
  File "/mnt/disk2/nestor/batvoice/dev/bv-etl-transcribe/env/lib/python3.9/site-packages/transformers/models/whisper/modeling_whisper.py", line 1834, in _extract_token_timestamps
    timestamps[batch_idx, 1:] = torch.tensor(jump_times)
RuntimeError: The expanded size of the tensor (105) must match the existing size (135) at non-singleton dimension 0.  Target sizes: [105].  Tensor sizes: [135]
```

### Expected behavior

When I change the parameter num_beams to 1, it works perfectly.

Is there any incompatibility between these two parameters? I've looked in the documentation and I haven't been able to find so.

The problem seems to be in this function (transformers/models/whisper/modeling_whisper.py):
```
    def _extract_token_timestamps(self, generate_outputs, alignment_heads, time_precision=0.02):
        """
        Calculates token-level timestamps using the encoder-decoder cross-attentions and dynamic time-warping (DTW) to
        map each output token to a position in the input audio.

        Returns:
            tensor containing the timestamps in seconds for each predicted token
        """
        # Create a list with `decoder_layers` elements, each a tensor of shape
        # (batch size, attention_heads, output length, input length).
        cross_attentions = []
        for i in range(self.config.decoder_layers):
            cross_attentions.append(torch.cat([x[i] for x in generate_outputs.cross_attentions], dim=2))

        # Select specific cross-attention layers and heads. This is a tensor
        # of shape (batch size, num selected, output length, input length).
        weights = torch.stack([cross_attentions[l][:, h] for l, h in alignment_heads])
        weights = weights.permute([1, 0, 2, 3])

        # Normalize and smoothen the weights.
        std, mean = torch.std_mean(weights, dim=-2, keepdim=True, unbiased=False)
        weights = (weights - mean) / std
        weights = _median_filter(weights, self.config.median_filter_width)

        # Average the different cross-attention heads.
        matrix = weights.mean(dim=1)

        timestamps = torch.zeros_like(generate_outputs.sequences, dtype=torch.float32)

        # Perform dynamic time warping on each element of the batch.
        for batch_idx in range(timestamps.shape[0]):
            text_indices, time_indices = _dynamic_time_warping(-matrix[batch_idx].double().cpu().numpy())
            jumps = np.pad(np.diff(text_indices), (1, 0), constant_values=1).astype(bool)
            jump_times = time_indices[jumps] * time_precision
            timestamps[batch_idx, 1:] = torch.tensor(jump_times)

        return timestamps
```

Thank you in advance.

Comments:
cc @sanchit-gandhi and @ylacombe 

Taking a look, thanks for taking the time to describe the issue! The code is not working though!

I believe this PR is taking care of this: #26699

(#28007 refers to the same issue)

Thanks!