Upload modeling_baichuan.py

modeling_baichuan.py

@@ -405,10 +405,10 @@ class BaichuanModel(BaichuanPreTrainedModel):
 
         if attention_mask is not None:
             if len(attention_mask.shape) == 2:
-                expanded_mask = attention_mask.to(alibi_mask.dtype)
+                expanded_mask = attention_mask.to(torch.float32)
                 expanded_mask = torch.tril(
-                    torch.gt(expanded_mask[:, :, None] * expanded_mask[:, None, :], 0)
-                ) * torch.eq(expanded_mask[:, :, None] - expanded_mask[:, None, :], 0)
+                    torch.gt(expanded_mask[:, :, None] * expanded_mask[:, None, :], 0).to(torch.float32)
+                ) * torch.eq(expanded_mask[:, :, None] - expanded_mask[:, None, :], 0).to(torch.float32)
             else:
                 expanded_mask = attention_mask
             bsz = inputs_embeds.size(0)