forked from Element-Research/dpnn
-
Notifications
You must be signed in to change notification settings - Fork 10
/
PrintSize.lua
36 lines (31 loc) · 848 Bytes
/
PrintSize.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
31
32
33
34
35
local PrintSize, parent = torch.class('nn.PrintSize', 'nn.Module')
function PrintSize:__init(prefix)
parent.__init(self)
self.prefix = prefix or "PrintSize"
end
function PrintSize:updateOutput(input)
self.output = input
local size
if torch.type(input) == 'table' then
size = input
elseif torch.type(input) == 'nil' then
size = 'missing size'
else
size = input:size()
end
print(self.prefix..":input\n", size)
return self.output
end
function PrintSize:updateGradInput(input, gradOutput)
local size
if torch.type(gradOutput) == 'table' then
size = gradOutput
elseif torch.type(gradOutput) == 'nil' then
size = 'missing size'
else
size = gradOutput:size()
end
print(self.prefix..":gradOutput\n", size)
self.gradInput = gradOutput
return self.gradInput
end