Fix Bugs

bert_layers.py

@@ -599,9 +599,9 @@ class BertModel(BertPreTrainedModel):
         device = input_ids.device if input_ids is not None else inputs_embeds.device
 
         if attention_mask is None:
-            attention_mask = torch.ones(input_shape, device=device)
+            attention_mask = torch.ones(input_shape, dtype=torch.long, device=device)
         if token_type_ids is None:
-            token_type_ids = torch.zeros(input_shape, device=device)
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
 
         embedding_output = self.embeddings(
             input_ids,