lib/torch/nn/functional.rb in torch-rb-0.1.1 vs lib/torch/nn/functional.rb in torch-rb-0.1.2

@@ -4,16 +4,30 @@
       class << self
         def relu(input)
           Torch.relu(input)
         end
 
-        def conv2d(input, weight, bias)
-          Torch.conv2d(input, weight, bias)
+        def conv2d(input, weight, bias, stride: 1, padding: 0)
+          # TODO pair stride and padding when needed
+          Torch.conv2d(input, weight, bias, stride, padding)
         end
 
+        def prelu(input, weight)
+          Torch.prelu(input, weight)
+        end
+
+        def leaky_relu(input, negative_slope = 0.01)
+          Torch.leaky_relu(input, negative_slope)
+        end
+
         def max_pool2d(input, kernel_size)
           kernel_size = [kernel_size, kernel_size] if kernel_size.is_a?(Integer)
           Torch.max_pool2d(input, kernel_size)
+        end
+
+        def avg_pool2d(input, kernel_size)
+          kernel_size = [kernel_size, kernel_size] if kernel_size.is_a?(Integer)
+          Torch.avg_pool2d(input, kernel_size)
         end
 
         def linear(input, weight, bias)
           Torch.linear(input, weight, bias)
         end