adding addtional weights parameter to ClassNLLCriterion, which is 1D tensor assigning different weight to each class

1 files changed, 16 insertions, 0 deletions

diff --git a/test/test.lua b/test/test.lua
index a0644fe..99b4bd5 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -512,6 +512,22 @@ function nntest.DistKLDivCriterion()
    criterionJacobianTest1D(cri, input, target)
 end
 
+function nntest.ClassNLLCriterion()
+   local numLabels = math.random(5,10)	
+   local input = torch.rand(numLabels)
+   local target = math.random(1,numLabels)	
+
+   -- default ClassNLLCriterion
+   local cri = nn.ClassNLLCriterion()
+   criterionJacobianTest1D(cri, input, target)
+
+   -- ClassNLLCriterion with weights
+   local weights = torch.rand(numLabels)
+   weights = weights / weights:sum()
+   cri = nn.ClassNLLCriterion(weights)
+   criterionJacobianTest1D(cri, input, target)
+end
+
 function nntest.LogSigmoid()
    local ini = math.random(10,20)
    local inj = math.random(10,20)