Added case in ConfusionMatrix, to handle NLL criterions.

1 files changed, 6 insertions, 1 deletions

diff --git a/ConfusionMatrix.lua b/ConfusionMatrix.lua
index 3486490..ca514f1 100644
--- a/ConfusionMatrix.lua
+++ b/ConfusionMatrix.lua
@@ -18,8 +18,13 @@ function ConfusionMatrix:add(prediction, target)
    if type(prediction) == 'number' then
       -- comparing numbers
       self.mat[target][prediction] = self.mat[target][prediction] + 1
+   elseif type(target) == 'number' then
+      -- prediction is a vector, then target assumed to be an index
+      local prediction_1d = torch.Tensor(prediction):resize(self.nclasses)
+      local _,prediction = lab.max(prediction_1d)
+      self.mat[target][prediction[1]] = self.mat[target][prediction[1]] + 1
    else
-      -- comparing vectors
+      -- both prediction and target are vectors
       local prediction_1d = torch.Tensor(prediction):resize(self.nclasses)
       local target_1d = torch.Tensor(target):resize(self.nclasses)
       local _,prediction = lab.max(prediction_1d)