MultiLabelMarginCriterion fixes for CUDA

1 files changed, 10 insertions, 0 deletions

diff --git a/Sum.lua b/Sum.lua
index 77d4fce..5d61c28 100644
--- a/Sum.lua
+++ b/Sum.lua
@@ -41,6 +41,11 @@ function Sum:updateGradInput(input, gradOutput)
     -- Instead, do a deepcopy
     local size      = input:size()
     size[dimension] = 1
+    if not gradOutput:isContiguous() then
+        self._gradOutput = self._gradOutput or gradOutput.new()
+        self._gradOutput:resizeAs(gradOutput):copy(gradOutput)
+        gradOutput = self._gradOutput
+    end
     gradOutput      = gradOutput:view(size)
     self.gradInput:resizeAs(input)
     self.gradInput:copy(gradOutput:expandAs(input))
@@ -49,3 +54,8 @@ function Sum:updateGradInput(input, gradOutput)
     end
     return self.gradInput
 end
+
+function Sum:clearState()
+    nn.utils.clear(self, '_gradOutput')
+    return parent.clearState(self)
+end