Don't bother filling a Tensor with zero right before we copy into it

2 files changed, 19 insertions, 2 deletions

diff --git a/gmodule.lua b/gmodule.lua
index 99698c0..674bcf9 100644
--- a/gmodule.lua
+++ b/gmodule.lua
@@ -11,8 +11,8 @@ local function getTotalGradOutput(node)
       node.data.gradOutputBuffer = node.data.gradOutputBuffer or nesting.cloneNested(gradOutput[1])
       local gobuff = node.data.gradOutputBuffer
       nesting.resizeNestedAs(gobuff, gradOutput[1])
-      nesting.fillNested(gobuff, 0)
-      for i=1,#gradOutput do
+      nesting.copyNested(gobuff, gradOutput[1])
+      for i=2,#gradOutput do
          nesting.addNestedTo(gobuff, gradOutput[i])
       end
       gradOutput = gobuff
diff --git a/nesting.lua b/nesting.lua
index 18899c1..e1ddd7b 100644
--- a/nesting.lua
+++ b/nesting.lua
@@ -51,6 +51,23 @@ function nesting.resizeNestedAs(output, input)
    end
 end
 
+-- Copies all tensors in the output.
+function nesting.copyNested(output, input)
+   if torch.isTensor(output) then
+      output:copy(input)
+   else
+      for key, child in pairs(input) do
+          nesting.copyNested(output[key], child)
+      end
+      -- Extra elements are removed from the output.
+      for key, child in pairs(output) do
+         if not input[key] then
+            output[key] = nil
+         end
+      end
+   end
+end
+
 -- Adds the input to the output.
 -- The input can contain nested tables.
 -- The output will contain the same nesting of tables.