Fix focal

kaidic · Nov 13, 2019 · 3193f05 · 3193f05
1 parent 6feb304
commit 3193f05
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/losses.py b/losses.py
@@ -18,7 +18,7 @@ def __init__(self, weight=None, gamma=0.):
         self.weight = weight
 
     def forward(self, input, target):
-        return focal_loss(F.cross_entropy(input, target, weight=self.weight), self.gamma)
+        return focal_loss(F.cross_entropy(input, target, reduction='none', weight=self.weight), self.gamma)
 
 class LDAMLoss(nn.Module):