Update modeling_quiet.py

modeling_quiet.py

@@ -136,7 +136,6 @@ def save_tokens_with_rewards_to_pdf(input_ids, token_rewards, tokenizer, output_
     c.save()
 
 
-# Copied from transformers.models.llama.modeling_llama._get_unpad_data
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
 
@@ -148,6 +147,8 @@ def _get_unpad_data(attention_mask):
     else:
         max_seqlen_in_batch = seqlens_in_batch.max().item()
 
+    # Ensure seqlens_in_batch has the correct shape before cumulative sum
+    seqlens_in_batch = seqlens_in_batch.view(-1)
     cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=attention_mask.device), seqlens_in_batch.cumsum(dim=0)])
 
     return indices, cu_seqlens, max_seqlen_in_batch