small fixes

torch · Jul 20, 2017 · 3371428 · 3371428
1 parent f1945f5
commit 3371428
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/ArgMax.lua b/ArgMax.lua
@@ -9,7 +9,7 @@ local ArgMax, parent = torch.class("nn.ArgMax", "nn.Module")
 function ArgMax:__init(dim, nInputDim, asLong)
    parent.__init(self)
    self.dim = dim or 1
-   self.nInputDim = nInputDim or 9999
+   self.nInputDim = nInputDim or 1
    self.asLong = (asLong == nil) and true or asLong
    if self.asLong then
       self.output = torch.LongTensor()
@@ -21,7 +21,7 @@ function ArgMax:updateOutput(input)
    self._indices = self._indices or
       (torch.type(input) == 'torch.CudaTensor' and (torch.CudaLongTensor and torch.CudaLongTensor() or torch.CudaTensor()) or torch.LongTensor())
    local dim = (input:dim() > self.nInputDim) and (self.dim + 1) or self.dim
-   
+
    torch.max(self._value, self._indices, input, dim)
    if input:dim() > 1 then
       local idx = self._indices:select(dim, 1)

diff --git a/Criterion.lua b/Criterion.lua
@@ -16,5 +16,5 @@ function Criterion:setZeroMask(zeroMask)
 end
 
 function Criterion:clearState()
-  return nn.utils.clear(self, 'gradInput')
+   return nn.utils.clear(self, 'gradInput')
 end