fixing *NLLCriterion for non mini-batch cases

torch · Oct 1, 2016 · 5ade793 · 5ade793
1 parent 688b7a7
commit 5ade793
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/ClassNLLCriterion.lua b/ClassNLLCriterion.lua
@@ -28,8 +28,10 @@ end
 
 function ClassNLLCriterion:updateOutput(input, target)
    if type(target) == 'number' then
-      if input:type() ~= 'torch.CudaTensor' then
-         self.target = self.target:long()
+      if input:type() == 'torch.CudaTensor' then
+          self.target = self.target:cudaLong()
+      else
+          self.target = self.target:long()
       end
       self.target[1] = target
    elseif target:type() == 'torch.CudaTensor' then
@@ -52,6 +54,11 @@ end
 
 function ClassNLLCriterion:updateGradInput(input, target)
    if type(target) == 'number' then
+      if input:type() == 'torch.CudaTensor' then
+          self.target = self.target:cudaLong()
+      else
+          self.target = self.target:long()
+      end
       self.target[1] = target
    elseif target:type() == 'torch.CudaTensor' then
       self.target = target:cudaLong()

diff --git a/SpatialClassNLLCriterion.lua b/SpatialClassNLLCriterion.lua
@@ -28,8 +28,10 @@ end
 
 function SpatialClassNLLCriterion:updateOutput(input, target)
    if type(target) == 'number' then
-      if input:type() ~= 'torch.CudaTensor' then
-         self.target = self.target:long()
+      if input:type() == 'torch.CudaTensor' then
+          self.target = self.target:cudaLong()
+      else
+          self.target = self.target:long()
       end
       self.target[1] = target
    elseif target:type() == 'torch.CudaTensor' then
@@ -52,6 +54,11 @@ end
 
 function SpatialClassNLLCriterion:updateGradInput(input, target)
    if type(target) == 'number' then
+      if input:type() == 'torch.CudaTensor' then
+          self.target = self.target:cudaLong()
+      else
+          self.target = self.target:long()
+      end
       self.target[1] = target
    elseif target:type() == 'torch.CudaTensor' then
       self.target = target:cudaLong()