Simpler code for DWConvClass (#4310)

* more simpler code for DWConvClass

more simpler code for DWConvClass

* remove DWConv function

* Replace DWConvClass with DWConv
This commit is contained in:
Yonghye Kwon authored and GitHub committed 2021-08-05 10:54:16 +02:00
1 parent f409d8e54f
commit e96c74b5a1
2 files changed
+3 -9

No files matched your search

+1 -1
View File
@@ -202,7 +202,7 @@ class Model(nn.Module):
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
LOGGER.info('Fusing layers... ')
for m in self.model.modules():
if isinstance(m, (Conv, DWConvClass)) and hasattr(m, 'bn'):
if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
delattr(m, 'bn') # remove batchnorm
m.forward = m.forward_fuse # update forward