fixbug of ops.conv

marsggbo · Apr 27, 2023 · bac0a1f · bac0a1f
1 parent 492dafb
commit bac0a1f
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/hyperbox/mutables/ops/conv.py b/hyperbox/mutables/ops/conv.py
@@ -360,6 +360,21 @@ def __init__(
         nn.Conv2d.__init__(self, **self.conv_kwargs)
         self.init_kernel_transform_matrix(kernel_size)
 
+    def forward(self, x):
+        out = None
+        if not self.is_search:
+            padding = self.padding
+            if self.auto_padding:
+                kernel_size = self.weight.shape[2:]
+                padding = []
+                for k in kernel_size:
+                    padding.append(k//2)
+                self.padding = padding
+            out = nn.Conv2d.forward(self, x)
+        else:
+            out = self.forward_conv(x)
+        return out
+
     def format_args(
         self,
         kernel_size: _size_2_t,
@@ -398,6 +413,21 @@ def __init__(
         nn.Conv3d.__init__(self, **self.conv_kwargs)
         self.init_kernel_transform_matrix(kernel_size)
 
+    def forward(self, x):
+        out = None
+        if not self.is_search:
+            padding = self.padding
+            if self.auto_padding:
+                kernel_size = self.weight.shape[2:]
+                padding = []
+                for k in kernel_size:
+                    padding.append(k//2)
+                self.padding = padding
+            out = nn.Conv3d.forward(self, x)
+        else:
+            out = self.forward_conv(x)
+        return out
+
     def format_args(
         self, 
         kernel_size: _size_3_t,