inplace HardTanh, subclass ReLU6

1 files changed, 9 insertions, 3 deletions

diff --git a/HardTanh.lua b/HardTanh.lua
index d3449a1..07cfc62 100644
--- a/HardTanh.lua
+++ b/HardTanh.lua
@@ -1,9 +1,13 @@
 local HardTanh, parent = torch.class('nn.HardTanh', 'nn.Module')
 
-function HardTanh:__init(min_value, max_value)
+function HardTanh:__init(min_value, max_value, inplace)
    parent.__init(self)
    self.min_val = min_value or -1
    self.max_val = max_value or 1
+   self.inplace = inplace or false
+   if (inplace and type(inplace) ~= 'boolean') then
+      error('in-place flag must be boolean')
+   end
    assert(self.max_val>self.min_val, 'max_value must be larger than min_value')
 end
 
@@ -14,7 +18,8 @@ function HardTanh:updateOutput(input)
       input:cdata(),
       self.output:cdata(),
       self.min_val,
-      self.max_val
+      self.max_val,
+      self.inplace or false
    )
    return self.output
 end
@@ -25,7 +30,8 @@ function HardTanh:updateGradInput(input, gradOutput)
       gradOutput:cdata(),
       self.gradInput:cdata(),
       self.min_val,
-      self.max_val
+      self.max_val,
+      self.inplace or false
    )
    return self.gradInput
 end