update

modeling_flexbert.py

@@ -1657,6 +1657,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
         )
+        print(hidden_states.shape)
 
         if self.compile_model:
             logits = self.compiled_lm_head(hidden_states)
@@ -1703,8 +1704,8 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
             # print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
             # print(new_logits.shape)
-            if new_logits.dim() == 1:
-                new_logits = new_logits.unsqueeze(0)
+            # if new_logits.dim() == 2:
+            #     new_logits = new_logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
                 logits=new_logits,
@@ -1713,8 +1714,9 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
             )
         else:
             print(f"Non-padding logits: {logits.shape}")
-            if logits.dim() == 1:
-                logits = logits.unsqueeze(0)
+            logits = logits.view(-1, logits.size(-1))
+            # if logits.dim() == 2:
+            #     logits = logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
                 logits=logits,