Merge pull request #2 from AdeelH/ignore_index

Fix ignore_index bug

focal_loss.py

@@ -35,6 +35,8 @@ def __init__(self,
         self.nll_loss = nn.NLLLoss(
             weight=alpha, reduction='none', ignore_index=ignore_index)
 
+        self.ignore_index = ignore_index
+
         if reduction in ('mean', 'sum', 'none'):
             self.reduction = reduction
         else:
@@ -49,6 +51,12 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor:
             # (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,)
             y = y.view(-1)
 
+        unignored_mask = y != self.ignore_index
+        y = y[unignored_mask]
+        if len(y) == 0:
+            return 0.
+        x = x[unignored_mask]
+
         # compute weighted cross entropy term: -alpha * log(pt)
         log_p = F.log_softmax(x, dim=-1)
         ce = self.nll_loss(log_p, y)