fix

modeling_flexbert.py

@@ -1643,7 +1643,14 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
 
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         if self.unpad_embeddings and (indices is None and cu_seqlens is None and max_seqlen is None):
-            batch_size, seq_len = input_ids.shape[:2]
+            if input_ids.dim() == 2:
+                batch_size, seq_len = input_ids.shape[:2]
+            elif input_ids.dim() >= 3:
+                batch_size, seq_len = input_ids.shape[:2]
+            else: # dim is 1
+                batch_size, seq_len = input_ids.shape[0], 1
+                input_ids = input_ids.unsqueeze(1)
+            
             if attention_mask is None:  # Create causal mask (lower triangular)
                 attention_mask = torch.tril(torch.ones(batch_size, seq_len, device=input_ids.device), diagonal=0)
             input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(