modeling

modeling_flexbert.py

@@ -1641,6 +1641,9 @@ 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 attention_mask is None:
+                # create all ones
+                attention_mask = torch.ones_like(input_ids)
             input_ids, indices, cu_seqlens, max_seqlen, position_ids, labels = self.unpad_inputs(
                 input_ids, attention_mask, position_ids, labels
             )