fix

modeling_flexbert.py

@@ -1644,7 +1644,7 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         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 causal mask (lower triangular)
-                attention_mask = torch.tril(torch.ones(batch_size, seq_len), diagonal=0)
+                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(
                 input_ids, attention_mask, position_ids, labels
             )