[Parallel,Multi]Criterion table inputs

1 files changed, 3 insertions, 3 deletions

diff --git a/ParallelCriterion.lua b/ParallelCriterion.lua
index 95bd6cc..84d4ee1 100644
--- a/ParallelCriterion.lua
+++ b/ParallelCriterion.lua
@@ -25,11 +25,11 @@ function ParallelCriterion:updateOutput(input, target)
 end
 
 function ParallelCriterion:updateGradInput(input, target)
+   self.gradInput = nn.utils.recursiveResizeAs(self.gradInput, input)
+   nn.utils.recursiveFill(self.gradInput, 0)
    for i,criterion in ipairs(self.criterions) do
       local target = self.repeatTarget and target or target[i]
-      self.gradInput[i] = self.gradInput[i] or input[i].new()
-      self.gradInput[i]:resizeAs(input[i]):zero()
-      self.gradInput[i]:add(self.weights[i], criterion:updateGradInput(input[i],target))
+      nn.utils.recursiveAdd(self.gradInput[i], self.weights[i], criterion:updateGradInput(input[i], target))
    end
    return self.gradInput
 end