PyTorch 1.6.0 compatibility updates

2020-07-28 14:48:33 -07:00 · 2020-07-28 14:48:33 -07:00 · 91af0401d8
parent 7f8471eaeb
commit 91af0401d8
1 changed files with 1 additions and 0 deletions
--- a/models/yolo.py
+++ b/models/yolo.py
@ -148,6 +148,7 @@ class Model(nn.Module):
        print('Fusing layers... ', end='')
        for m in self.model.modules():
            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.bn = None  # remove batchnorm
                m.forward = m.fuseforward  # update forward