PyTorch 1.6.0 compatibility updates
Browse files- models/yolo.py +1 -0
models/yolo.py
CHANGED
@@ -148,6 +148,7 @@ class Model(nn.Module):
|
|
148 |
print('Fusing layers... ', end='')
|
149 |
for m in self.model.modules():
|
150 |
if type(m) is Conv:
|
|
|
151 |
m.conv = torch_utils.fuse_conv_and_bn(m.conv, m.bn) # update conv
|
152 |
m.bn = None # remove batchnorm
|
153 |
m.forward = m.fuseforward # update forward
|
|
|
148 |
print('Fusing layers... ', end='')
|
149 |
for m in self.model.modules():
|
150 |
if type(m) is Conv:
|
151 |
+
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
|
152 |
m.conv = torch_utils.fuse_conv_and_bn(m.conv, m.bn) # update conv
|
153 |
m.bn = None # remove batchnorm
|
154 |
m.forward = m.fuseforward # update forward
|