model.fuse() fix for export.py (#827)

This commit is contained in:
Glenn Jocher committed 2020-08-24 19:27:54 -07:00
1 parent 5f07782f56
commit a8751e50de
2 files changed
+4 -2

No files matched your search

+1 -1
View File
@@ -163,7 +163,7 @@ class Model(nn.Module):
if type(m) is Conv:
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
m.bn = None # remove batchnorm
delattr(m, 'bn') # remove batchnorm
m.forward = m.fuseforward # update forward
self.info()
return self