Merge pull request #22 from wkcn/fix_conv_bn_fuse

[fix] Fix Conv2d_BN fuse bug when groups > 1
facebookresearch · Mar 25, 2022 · d000b74 · d000b74
2 parents d373b82 + 3bc3ec2
commit d000b74
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/levit.py b/levit.py
@@ -94,7 +94,7 @@ def fuse(self):
         w = c.weight * w[:, None, None, None]
         b = bn.bias - bn.running_mean * bn.weight / \
             (bn.running_var + bn.eps)**0.5
-        m = torch.nn.Conv2d(w.size(1), w.size(
+        m = torch.nn.Conv2d(w.size(1) * self.c.groups, w.size(
             0), w.shape[2:], stride=self.c.stride, padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups)
         m.weight.data.copy_(w)
         m.bias.data.copy_(b)

diff --git a/levit_c.py b/levit_c.py
@@ -94,7 +94,7 @@ def fuse(self):
         w = c.weight * w[:, None, None, None]
         b = bn.bias - bn.running_mean * bn.weight / \
             (bn.running_var + bn.eps)**0.5
-        m = torch.nn.Conv2d(w.size(1), w.size(
+        m = torch.nn.Conv2d(w.size(1) * self.c.groups, w.size(
             0), w.shape[2:], stride=self.c.stride, padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups)
         m.weight.data.copy_(w)
         m.bias.data.copy_(b)