Linear:updateGradInput avoids NaN and inf

The resize in Linear:updateGradInput can introduce NaN and inf into the
gradients. The resize itself leaves garbage in the gradInput tensor.
For normal numbers the subsequent addmm/addmv will clear the garbage.
However, if gradInput contains either nan or inf after the resize, the
multiply will result in nan instead of the desired result.

1 files changed, 5 insertions, 2 deletions

diff --git a/Linear.lua b/Linear.lua
index cc6da4e..2e6635c 100644
--- a/Linear.lua
+++ b/Linear.lua
@@ -52,11 +52,14 @@ end
 function Linear:updateGradInput(input, gradOutput)
    if self.gradInput then
 
+      local nElement = self.gradInput:nElement()
+      self.gradInput:resizeAs(input)
+      if self.gradInput:nElement() ~= nElement then
+         self.gradInput:zero()
+      end
       if input:dim() == 1 then
-         self.gradInput:resizeAs(input)
          self.gradInput:addmv(0, 1, self.weight:t(), gradOutput)
       elseif input:dim() == 2 then
-         self.gradInput:resizeAs(input)
          self.gradInput:addmm(0, 1, gradOutput, self.weight)
       end