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
36

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