Gemma2: fix FA2 generation (#32553)

fix FA2
huggingface · Aug 20, 2024 · 679e235 · 679e235
1 parent 984bc11
commit 679e235
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/transformers/models/gemma2/modeling_gemma2.py b/src/transformers/models/gemma2/modeling_gemma2.py
@@ -1083,9 +1083,9 @@ def prepare_inputs_for_generation(
             model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format)}
 
         if isinstance(past_key_values, HybridCache) and attention_mask.ndim == 2:
-            if inputs_embeds is not None:
-                batch_size, sequence_length = inputs_embeds.shape
-                device = inputs_embeds.device
+            if model_inputs["inputs_embeds"] is not None:
+                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
+                device = model_inputs["inputs_embeds"].device
             else:
                 batch_size, sequence_length = input_ids.shape
                 device = input_ids.device