Update modeling_quiet.py

modeling_quiet.py

@@ -139,14 +139,18 @@ def save_tokens_with_rewards_to_pdf(input_ids, token_rewards, tokenizer, output_
 # 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)
+
     indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
-    return (
-        indices,
-        cu_seqlens,
-        max_seqlen_in_batch,
-    )
+
+    # Handle the case when seqlens_in_batch is empty
+    if seqlens_in_batch.numel() == 0:
+        max_seqlen_in_batch = 0
+    else:
+        max_seqlen_in_batch = seqlens_in_batch.max().item()
+
+    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
 
 
 # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Quiet