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