Update extra/nn/Sum.lua

make sure that self.output is a tensor

1 files changed, 3 insertions, 0 deletions

diff --git a/Sum.lua b/Sum.lua
index b068e25..fa22309 100644
--- a/Sum.lua
+++ b/Sum.lua
@@ -7,6 +7,9 @@ function Sum:__init(dimension)
 end
 
 function Sum:updateOutput(input)
+   if type(self.output) == 'number' then
+      self.output = input.new()
+   end
    input.torch.sum(self.output, input, self.dimension)
    self.output = self.output:select(self.dimension, 1)
    return self.output