🔹 Fix: Miscalculated mask shape in comments (#2925)

Author: linkedlist771

trl/trainer/grpo_trainer.py

@@ -739,7 +739,7 @@ def _generate_and_score_completions(
         completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int()
 
         # Concatenate prompt_mask with completion_mask for logit computation
-        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B*G, P+C)
+        attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)  # (B, P+C)
 
         logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens