update

modeling_flexbert.py

@@ -1703,7 +1703,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
 
         if self.unpad_embeddings:
             # reshape to batch size
-            logits = logits.view(batch_size, -1, self.vocab_size)
+            logits = logits.view(-1, self.vocab_size)
 
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")