added

modeling_flexbert.py

@@ -1536,7 +1536,7 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
         # Initialize weights and apply final processing
         self._init_weights(reset_params=False)
 
-    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
+[]    def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
         # Handle the XOR condition
         assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"
         
@@ -1556,7 +1556,7 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
             
             if not self.config.tie_word_embeddings:
                 init_weights(self.config, self.decoder, self.config.hidden_size, type_of_module=ModuleType.final_out)
-                
+
     @classmethod
     def from_composer(
         cls,
@@ -1702,13 +1702,19 @@ class FlexBertForCasualLM(FlexBertPreTrainedModel):
             )
 
         if self.pad_logits:
+            # Reshape logits to 3D if needed
+            new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
+            if len(new_logits.shape) == 2:
+                new_logits = new_logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
-                logits=self.pad_inputs(logits, indices, batch_size, seq_len)[0],
+                logits=new_logits,
                 hidden_states=None,
                 attentions=None,
             )
         else:
+            if len(logits.shape) == 2:
+                logits = logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
                 logits=logits,