blob: 3178e9272ab44afd2acac055a2cb870defe8abba (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
|
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
|