zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

PyTorch 1.6.0 compatibility updates

This commit is contained in:
Glenn Jocher 2020-07-28 14:48:33 -07:00 committed by GitHub
parent 7f8471eaeb
commit 91af0401d8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
models/yolo.py
View File

@ -148,6 +148,7 @@ class Model(nn.Module):
print('Fusing layers... ', end='') print('Fusing layers... ', end='')
for m in self.model.modules(): for m in self.model.modules():
if type(m) is Conv: if type(m) is Conv:
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
m.conv = torch_utils.fuse_conv_and_bn(m.conv, m.bn) # update conv m.conv = torch_utils.fuse_conv_and_bn(m.conv, m.bn) # update conv
m.bn = None # remove batchnorm m.bn = None # remove batchnorm
m.forward = m.fuseforward # update forward m.forward = m.fuseforward # update forward