nn.Select accepts negative indices

1 files changed, 4 insertions, 2 deletions

diff --git a/Select.lua b/Select.lua
index acf8e06..fccdf32 100644
--- a/Select.lua
+++ b/Select.lua
@@ -7,14 +7,16 @@ function Select:__init(dimension,index)
 end
 
 function Select:updateOutput(input)
-   local output = input:select(self.dimension,self.index);
+   local index = self.index < 0 and input:size(self.dimension) + self.index + 1 or self.index
+   local output = input:select(self.dimension, index);
    self.output:resizeAs(output)
    return self.output:copy(output)
 end
 
 function Select:updateGradInput(input, gradOutput)
+   local index = self.index < 0 and input:size(self.dimension) + self.index + 1 or self.index
    self.gradInput:resizeAs(input)  
    self.gradInput:zero()
-   self.gradInput:select(self.dimension,self.index):copy(gradOutput) 
+   self.gradInput:select(self.dimension,index):copy(gradOutput) 
    return self.gradInput
 end