ops wrong order

Author: manueldeprada

src/transformers/models/gemma3n/modeling_gemma3n.py

@@ -1355,10 +1355,10 @@ def forward(
                 "cache_position": cache_position,
                 "sliding_window": self.sliding_window,
             }
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
             if self.store_full_length_kv:
                 past_key_values.layers[self.layer_idx].keys_to_share = key_states
                 past_key_values.layers[self.layer_idx].values_to_share = value_states
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":

src/transformers/models/gemma3n/modular_gemma3n.py

@@ -1806,10 +1806,10 @@ def forward(
                 "cache_position": cache_position,
                 "sliding_window": self.sliding_window,
             }
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
             if self.store_full_length_kv:
                 past_key_values.layers[self.layer_idx].keys_to_share = key_states
                 past_key_values.layers[self.layer_idx].values_to_share = value_states
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":