Fix NaN in Gemma3/EmbeddingGemma when batching mixed-length sequences…

src/transformers/masking_utils.py

@@ -610,6 +610,14 @@ def eager_mask(
         min_dtype = torch.finfo(dtype).min
         # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type)
         mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype)
+        # Prevent NaN from softmax([-inf, ..., -inf]) for rows where all keys are masked. This occurs when a
+        # padding query position's sliding window falls entirely within the padding region (e.g. a short sequence
+        # padded to match a much longer one in a batch). Unlike torch.nn.functional.scaled_dot_product_attention
+        # (which handles this natively since torch>=2.5), the manual softmax in eager attention forward does not.
+        # For such rows we allow all keys (bias=0), producing a valid uniform distribution. These rows always
+        # correspond to padding positions and are excluded from final pooling, so the result is unused.
+        all_masked_rows = (mask == min_dtype).all(dim=-1, keepdim=True)
+        mask = torch.where(all_masked_rows, torch.zeros_like(mask), mask)
     return mask