Fix nan issue in nemo1 / neva (#11629)

* potential fix Signed-off-by: yaoyu-33 <[email protected]> * Apply isort and black reformatting Signed-off-by: yaoyu-33 <[email protected]> --------- Signed-off-by: yaoyu-33 <[email protected]> Signed-off-by: yaoyu-33 <[email protected]> Co-authored-by: yaoyu-33 <[email protected]>
NVIDIA · Jan 7, 2025 · 79363ac · 79363ac
1 parent 8f3485e
commit 79363ac
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nemo/collections/multimodal/models/multimodal_llm/neva/neva_model.py b/nemo/collections/multimodal/models/multimodal_llm/neva/neva_model.py
@@ -1189,8 +1189,10 @@ def test_epoch_end(self, outputs):
     def loss_func(self, loss_mask, output_tensor):
         losses = output_tensor.float()
         loss_mask = loss_mask.view(-1).float()
-        # TODO: add nemo version here
-        loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()  # sequence level nll
+        valid_tokens = loss_mask.sum()
+        if valid_tokens < 0.5:  # no valid tokens
+            valid_tokens += 1.0
+        loss = torch.sum(losses.view(-1) * loss_mask) / valid_tokens  # sequence level nll
         return loss
 
     def setup(self, stage=None):