fuse() bug fix

2020-10-11 00:42:54 +02:00 · 2020-10-11 00:42:54 +02:00 · c2403d785c
parent 5fac5ad165
commit c2403d785c
1 changed files with 1 additions and 1 deletions
--- a/models/yolo.py
+++ b/models/yolo.py
@ -162,7 +162,7 @@ class Model(nn.Module):
    def fuse(self):  # fuse model Conv2d() + BatchNorm2d() layers
        print('Fusing layers... ')
        for m in self.model.modules():
-            if type(m) is Conv and hasattr(Conv, 'bn'):
+            if type(m) is Conv and hasattr(m, 'bn'):
                m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
                m.conv = fuse_conv_and_bn(m.conv, m.bn)  # update conv
                delattr(m, 'bn')  # remove batchnorm