add

modeling_flexbert.py

@@ -1529,6 +1529,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         self.unpad_embeddings = config.unpad_embeddings
         self.pad_logits = config.pad_logits
         self.compile_model = config.compile_model
+        self.vocab_size = config.vocab_size
         # self.masked_prediction = config.masked_prediction
 
         # Initialize weights and apply final processing
@@ -1702,7 +1703,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
 
         if self.unpad_embeddings:
             # reshape to batch size
-            logits = logits.view(batch_size, -1, logits.size(-1))
+            logits = logits.view(batch_size, -1, self.vocab_size)
 
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")