contrib/torch/nn/Power.lua


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

local Power, parent = torch.class('nn.Power','nn.Module')

function Power:__init(p)
   parent.__init(self)
   self.pow = p
   if not p then
      error('nn.Power(power)')
   end
end

function Power:updateOutput(input)
   self.output:resizeAs(input):copy(input)
   self.output:pow(self.pow)
   return self.output
end

function Power:updateGradInput(input, gradOutput)
   self.gradInput:resizeAs(input):copy(input)
   self.gradInput:pow(self.pow - 1)
   self.gradInput:cmul(gradOutput):mul(self.pow)
   return self.gradInput
end