modeling

modeling_flexbert.py

@@ -1700,6 +1700,10 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels.view(-1)
             )
 
+        if self.unpad_embeddings:
+            # reshape to batch size
+            logits = logits.view(batch_size, -1, logits.size(-1))
+
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]