Avoid Nan Losses (#2001)

Fixes small error where nan losses could appear if whole input is
masked.

model/model_training/losses.py

@@ -4,7 +4,7 @@
 
 
 class CrossEntropyLoss(nn.CrossEntropyLoss):
-    def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean"):
+    def __init__(self, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="none"):
         super().__init__(weight, size_average, ignore_index, reduce, reduction)
 
     def forward(self, input, target, mask=None):
@@ -14,7 +14,12 @@ def forward(self, input, target, mask=None):
             target = target.view(-1)
             input = input[mask]
             target = target[mask]
-        return super().forward(input, target)
+
+        size = target.numel()
+
+        loss = super().forward(input, target)
+
+        return loss.sum() / (size + 1e-8)
 
 
 class PolyLoss(nn.Module):