2018-08-29 08:47:45 +12:00
|
|
|
local Print, parent = torch.class('w2nn.Print','nn.Module')
|
|
|
|
|
2018-10-14 05:20:03 +13:00
|
|
|
function Print:__init(id)
|
2018-08-29 08:47:45 +12:00
|
|
|
parent.__init(self)
|
2018-10-14 05:20:03 +13:00
|
|
|
self.id = id
|
2018-08-29 08:47:45 +12:00
|
|
|
end
|
|
|
|
function Print:updateOutput(input)
|
2018-10-14 05:20:03 +13:00
|
|
|
print("----", self.id)
|
2018-08-29 08:47:45 +12:00
|
|
|
print(input:size())
|
|
|
|
self.output:resizeAs(input)
|
|
|
|
self.output:copy(input)
|
|
|
|
return self.output
|
|
|
|
end
|
|
|
|
function Print:updateGradInput(input, gradOutput)
|
|
|
|
self.gradInput:resizeAs(GradOutput)
|
|
|
|
return self.gradInput
|
|
|
|
end
|