State variable pointers should never change.

see: https://github.com/torch/nn/blob/master/doc/module.md#state-variables

1 files changed, 2 insertions, 2 deletions

diff --git a/Identity.lua b/Identity.lua
index 088cc34..18a3c9b 100644
--- a/Identity.lua
+++ b/Identity.lua
@@ -1,12 +1,12 @@
 local Identity, _ = torch.class('nn.Identity', 'nn.Module')
 
 function Identity:updateOutput(input)
-   self.output = input
+   self.output:set(input:storage(), input:storageOffset(), input:size(), input:stride())
    return self.output
 end
 
 
 function Identity:updateGradInput(input, gradOutput)
-   self.gradInput = gradOutput
+   self.gradInput:set(gradOutput:storage(), gradOutput:storageOffset(), gradOutput:size(), gradOutput:stride())
    return self.gradInput
 end