diff options
| -rw-r--r-- | Criterion.lua | 6 | ||||
| -rw-r--r-- | Module.lua | 12 |
2 files changed, 17 insertions, 1 deletions
diff --git a/Criterion.lua b/Criterion.lua index 6513414..dec514c 100644 --- a/Criterion.lua +++ b/Criterion.lua @@ -49,3 +49,9 @@ end function Criterion:cuda() return self:type('torch.CudaTensor') end + +function Criterion:__call__(input, target) + self.output = self:forward(input, target) + self.gradInput = self:backward(input, target) + return self.output, self.gradInput +end @@ -22,7 +22,7 @@ function Module:updateOutput(input) end function Module:forward(input) - return self:updateOutput(input, target) + return self:updateOutput(input) end function Module:backward(input, gradOutput) @@ -212,3 +212,13 @@ function Module:getParameters() -- return new flat vector that contains all discrete parameters return flatParameters, flatGradParameters end + +function Module:__call__(input, gradOutput) + self:forward(input) + if gradOutput then + self:backward(input, gradOutput) + return self.output, self.gradInput + else + return self.output + end +end |