Commit
·
2f77cf3
1
Parent(s):
89655a8
.fuse() additional error checking
Browse files- models/yolo.py +1 -1
models/yolo.py
CHANGED
@@ -160,7 +160,7 @@ class Model(nn.Module):
|
|
160 |
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
161 |
print('Fusing layers... ')
|
162 |
for m in self.model.modules():
|
163 |
-
if type(m) is Conv:
|
164 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
165 |
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
166 |
delattr(m, 'bn') # remove batchnorm
|
|
|
160 |
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
|
161 |
print('Fusing layers... ')
|
162 |
for m in self.model.modules():
|
163 |
+
if type(m) is Conv and hasattr(Conv, 'bn'):
|
164 |
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
165 |
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
|
166 |
delattr(m, 'bn') # remove batchnorm
|