FunctionWrapper.lua

require 'nnx'
local FunctionWrapper, parent = torch.class('nn.FunctionWrapper', 'nn.Module')

local help_desc = [[
      Dummy module that takes a forward and a backward function as argument.
]]

function FunctionWrapper:__init(init, updateOutput, updateGradInput)
   init(self)
   self.fn_updateOutput = updateOutput
   self.fn_updateGradInput = updateGradInput
end

function FunctionWrapper:updateOutput(input)
   self.output = self.fn_updateOutput(self, input)
   return self.output
end

function FunctionWrapper:updateGradInput(input, gradOutput)
   self.gradInput = self.fn_updateGradInput(self, input, gradOutput)
   return self.gradInput
end