Merge pull request #7 from AdeelH/AdeelH-patch-1

Ensure forward() always returns a tensor
AdeelH · Jun 18, 2022 · 6359c1e · 6359c1e
2 parents f15f440 + 00c99ad
commit 6359c1e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/focal_loss.py b/focal_loss.py
@@ -66,7 +66,7 @@ def forward(self, x: Tensor, y: Tensor) -> Tensor:
         unignored_mask = y != self.ignore_index
         y = y[unignored_mask]
         if len(y) == 0:
-            return 0.
+            return torch.tensor(0.)
         x = x[unignored_mask]
 
         # compute weighted cross entropy term: -alpha * log(pt)