fuse() bug fix
Browse files- models/yolo.py +1 -1
models/yolo.py
CHANGED
@@ -162,7 +162,7 @@ class Model(nn.Module):
|
|
162 |
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
163 |
print('Fusing layers... ')
|
164 |
for m in self.model.modules():
|
165 |
-
if type(m) is Conv and hasattr(
|
166 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
167 |
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
168 |
delattr(m, 'bn') # remove batchnorm
|
|
|
162 |
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
163 |
print('Fusing layers... ')
|
164 |
for m in self.model.modules():
|
165 |
+
if type(m) is Conv and hasattr(m, 'bn'):
|
166 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
167 |
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
168 |
delattr(m, 'bn') # remove batchnorm
|