Fixing bug in Megatron BERT when loss mask is all zeros

nemo/collections/nlp/models/language_modeling/megatron_bert_model.py

@@ -358,7 +358,12 @@ def loss_func(self, loss_mask, sentence_order, output_tensor):
 
         lm_loss_ = lm_loss_.float()
         loss_mask = loss_mask.float()
-        lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+
+        # Sometimes when the number of tokens is very small, none of the tokens get masked for prediction. In that case loss mask is all zeros
+        if loss_mask.sum() == 0:
+            lm_loss = torch.sum(lm_loss_.view(-1)) * 0.0
+        else:
+            lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
 
         if sop_logits is not None:
             sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(), sentence_order.view(-1), ignore_index=-1)