lib/torch/nn/relu.rb in torch-rb-0.1.3 vs lib/torch/nn/relu.rb in torch-rb-0.1.4

@@ -1,13 +1,18 @@
 module Torch
   module NN
     class ReLU < Module
-      def initialize #(inplace: false)
-        # @inplace = inplace
+      def initialize(inplace: false)
+        super()
+        @inplace = inplace
       end
 
       def forward(input)
-        F.relu(input) #, inplace: @inplace)
+        F.relu(input, inplace: @inplace)
+      end
+
+      def extra_inspect
+        @inplace ? "inplace: true" : ""
       end
     end
   end
 end