Fix fuse (#440)

Fix fuse in models.py
2019-08-09 18:44:47 +08:00 · 2019-08-09 18:44:47 +08:00 · d41b444a15
parent fdd5afa229
commit d41b444a15
1 changed files with 8 additions and 7 deletions
--- a/models.py
+++ b/models.py
@ -213,6 +213,7 @@ class Darknet(nn.Module):
        # Fuse Conv2d + BatchNorm2d layers throughout model
        fused_list = nn.ModuleList()
        for a in list(self.children())[0]:
+            if isinstance(a, nn.Sequential):
                for i, b in enumerate(a):
                    if isinstance(b, nn.modules.batchnorm.BatchNorm2d):
                        # fuse this bn layer with the previous conv2d layer