fix BertForMaskedLM

modeling_bert.py

@@ -752,18 +752,18 @@ class BertForMaskedLM(BertPreTrainedModel):
         prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
 
         if (
-            self.dense_seq_output and labels is not None
+                self.dense_seq_output and labels is not None
         ):  # prediction_scores are already flattened
             masked_lm_loss = self.mlm_loss(
                 prediction_scores, labels.flatten()[masked_token_idx]
             ).float()
-
-        assert labels is not None
-
-        masked_lm_loss = self.mlm_loss(
-            rearrange(prediction_scores, "... v -> (...) v"),
-            rearrange(labels, "... -> (...)"),
-        ).float()
+        elif labels is not None:
+            masked_lm_loss = self.mlm_loss(
+                rearrange(prediction_scores, "... v -> (...) v"),
+                rearrange(labels, "... -> (...)"),
+            ).float()
+        else:
+            raise ValueError('MLM labels must not be None')
 
         return BertForPreTrainingOutput(
             loss=masked_lm_loss,