Update legacy ClassNLLCriterion to add ignore_index.

760chong · Aug 1, 2017 · 9c1e9d8 · 9c1e9d8
1 parent 61c873c
commit 9c1e9d8
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/torch/legacy/nn/ClassNLLCriterion.py b/torch/legacy/nn/ClassNLLCriterion.py
@@ -4,9 +4,10 @@
 
 class ClassNLLCriterion(Criterion):
 
-    def __init__(self, weights=None, sizeAverage=True):
+    def __init__(self, weights=None, sizeAverage=True, ignore_index=-100):
         super(ClassNLLCriterion, self).__init__()
         self.sizeAverage = sizeAverage
+        self.ignore_index = ignore_index
 
         if weights is not None:
             assert weights.dim() == 1
@@ -25,7 +26,7 @@ def updateOutput(self, input, target):
             self.sizeAverage,
             self.weights,
             self.total_weight_tensor,
-            -100
+            self.ignore_index
         )
         self.output = self.output_tensor[0]
         return self.output
@@ -42,7 +43,7 @@ def updateGradInput(self, input, target):
             self.sizeAverage,
             self.weights,
             self.total_weight_tensor,
-            -100
+            self.ignore_index
         )
 
         return self.gradInput