-
Notifications
You must be signed in to change notification settings - Fork 966
/
Collapse.lua
30 lines (26 loc) · 743 Bytes
/
Collapse.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
local Collapse, parent = torch.class('nn.Collapse', 'nn.Module')
-- collapses non-batch dims
function Collapse:__init(nInputDim)
parent.__init(self)
self.nInputDim = nInputDim
end
function Collapse:updateOutput(input)
if not input:isContiguous() then
self._input = self._input or input.new()
self._input:resize(input:size()):copy(input)
input = self._input
end
if input:dim() > self.nInputDim then
self.output:view(input,input:size(1),-1)
else
self.output:view(input,-1)
end
return self.output
end
function Collapse:updateGradInput(input, gradOutput)
self.gradInput:view(gradOutput, input:size())
return self.gradInput
end
function Collapse:clearState()
self._input = nil
end