updated MLM predictions

modeling_gptbert.py

@@ -778,6 +778,10 @@ class GptBertForMaskedLM(GptBertModel):
             subword_prediction_flatten = subword_prediction[:, :-1].flatten(0, 1)
             masked_lm_loss = F.cross_entropy(subword_prediction_flatten, labels_flatten)
 
+        bos_logits = torch.zeros(subword_prediction.size(0), 1, self.config.vocab_size, dtype=subword_prediction.dtype, device=subword_prediction.device)
+        bos_logits[:, :, self.config.bos_token_id] = 1.0
+        subword_prediction = torch.cat([bos_logits, subword_prediction[:, :-1]], dim=1)
+
         if not return_dict:
             output = (
                 subword_prediction,