Repository: huggingface/transformers
Pull request: Fixes the BC

# What does this PR do?
Fix #39558

The bug fixe now shows:
```python
Number of returned hidden states: 17
hidden state  0 values (slice): tensor([0.0045, 0.0166, 0.0210])
hidden state  1 values (slice): tensor([-0.0010,  0.0339,  0.0018])
hidden state  2 values (slice): tensor([ 0.2473, -0.5012,  1.5765])
hidden state  3 values (slice): tensor([ 0.2444, -0.5238,  1.5824])
hidden state  4 values (slice): tensor([ 0.2463, -0.5339,  1.5775])
hidden state  5 values (slice): tensor([ 0.2797, -0.5208,  1.6037])
hidden state  6 values (slice): tensor([ 0.2424, -0.4786,  1.6114])
hidden state  7 values (slice): tensor([ 0.2291, -0.4361,  1.5910])
hidden state  8 values (slice): tensor([ 0.2474, -0.3626,  1.5842])
hidden state  9 values (slice): tensor([ 0.2487, -0.3277,  1.4880])
hidden state 10 values (slice): tensor([ 0.1552, -0.2746,  1.4607])
hidden state 11 values (slice): tensor([ 0.1235, -0.2442,  1.3794])
hidden state 12 values (slice): tensor([ 0.1107, -0.2346,  1.3300])
hidden state 13 values (slice): tensor([ 0.0789, -0.2590,  1.3076])
hidden state 14 values (slice): tensor([ 0.3292, -0.3711,  1.1052])
hidden state 15 values (slice): tensor([ 0.4628, -0.6839,  1.1164])
hidden state 16 values (slice): tensor([ 0.2918, -0.4089,  1.8382])
```

This will not work for gemma, but the effort for gemma is to record the hidden states before the forward, a hack but can be done if this break anyone's workflow:
```python 
        def make_capture_wrapper(module, orig_forward, key, index):
            @wraps(orig_forward)
            def wrapped_forward(*args, **kwargs):
                if "hidden_states" in collected_outputs:
                    collected_outputs[key] += args[0]
                output = orig_forward(*args, **kwargs)
                if not isinstance(output, tuple):
                    collected_outputs[key] += (output,)
                elif output[index] is not None:
                    collected_outputs[key] += (output[index],)
                return output

            return wrapped_forward
```
something like that 😉 


Changed files:
filename: src/transformers/utils/generic.py
patch: @@ -1032,6 +1032,8 @@ def wrapper(self, *args, **kwargs):
         def make_capture_wrapper(module, orig_forward, key, index):
             @wraps(orig_forward)
             def wrapped_forward(*args, **kwargs):
+                if key == "hidden_states" and len(collected_outputs[key]) == 0:
+                    collected_outputs[key] += (args[0],)
                 output = orig_forward(*args, **kwargs)
                 if not isinstance(output, tuple):
                     collected_outputs[key] += (output,)
@@ -1065,18 +1067,19 @@ def wrapped_forward(*args, **kwargs):
                         monkey_patched_layers.append((module, original_forward))
 
         outputs = func(self, *args, **kwargs)
-
         # Restore original forward methods
         for module, original_forward in monkey_patched_layers:
             module.forward = original_forward
 
         # Inject collected outputs into model output
         for key in collected_outputs:
             if key == "hidden_states":
+                collected_outputs[key] = collected_outputs[key][:-1]
                 if hasattr(outputs, "vision_hidden_states"):
                     collected_outputs[key] += (outputs.vision_hidden_states,)
                 elif hasattr(outputs, "last_hidden_state"):
                     collected_outputs[key] += (outputs.last_hidden_state,)
+
                 outputs[key] = collected_outputs[key]
             elif key == "attentions":
                 if isinstance(capture_flags[key], list) and len(capture_flags[key]) == 2: