glenn-jocher commited on
Commit
2f77cf3
·
1 Parent(s): 89655a8

.fuse() additional error checking

Browse files
Files changed (1) hide show
  1. 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