Retain nn.Identity state pointers if they are tensors

1 files changed, 21 insertions, 3 deletions

diff --git a/Identity.lua b/Identity.lua
index 088cc34..3d73f81 100644
--- a/Identity.lua
+++ b/Identity.lua
@@ -1,12 +1,30 @@
-local Identity, _ = torch.class('nn.Identity', 'nn.Module')
+local Identity, parent = torch.class('nn.Identity', 'nn.Module')
+
+function Identity:__init()
+   parent.__init(self)
+   self.tensorOutput = torch.Tensor{}
+   self.output = self.tensorOutput
+   self.tensorGradInput = torch.Tensor{}
+   self.gradInput = self.tensorGradInput
+end
 
 function Identity:updateOutput(input)
-   self.output = input
+   if torch.isTensor(input) then
+      self.tensorOutput:set(input)
+      self.output = self.tensorOutput
+   else
+      self.output = input
+   end
    return self.output
 end
 
 
 function Identity:updateGradInput(input, gradOutput)
-   self.gradInput = gradOutput
+   if torch.isTensor(gradOutput) then
+      self.tensorGradInput:set(gradOutput)
+      self.gradInput = self.tensorGradInput
+   else
+      self.gradInput = gradOutput
+   end
    return self.gradInput
 end