Now you can use power on inputs with 0s. (there is no longer division by 0)

1 files changed, 24 insertions, 21 deletions

diff --git a/Power.lua b/Power.lua
index 8052b3f..76421c9 100644
--- a/Power.lua
+++ b/Power.lua
@@ -1,21 +1,24 @@
-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(gradOutput)
-   self.gradInput:cmul(self.output):cdiv(input):mul(self.pow)
-   return self.gradInput
-end
+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.buffer = self.buffer or input.new()

+   self.buffer:resizeAs(input):copy(input)

+   self.buffer:pow(self.pow - 1)

+   self.gradInput:resizeAs(input):copy(gradOutput)

+   self.gradInput:cmul(self.buffer):mul(self.pow)

+   return self.gradInput

+end
\ No newline at end of file