fix(gemma4): prevent attention mask offset overflow (#31672)

Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
2026-07-18 23:24:03 -07:00
parent b8ec544946
commit 609fe1c0d1
1 changed files with 1 additions and 1 deletions
@@ -335,7 +335,7 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):

        bidirectional_attn_masks_list = []
        bidirectional_attn_mask_indptr = torch.zeros(
-            forward_batch.batch_size + 1, dtype=torch.int32, device=input_ids.device
+            forward_batch.batch_size + 1, dtype=torch.int64, device=input_ids.device
        )

        split_images = []