fuse() bug fix

This commit is contained in:
Glenn Jocher authored and GitHub committed 2020-10-11 00:42:54 +02:00
1 parent 5fac5ad165
commit c2403d785c
1 file changed
+1 -1
+1 -1
View File
@@ -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