glenn-jocher commited on
Commit
91af040
·
unverified ·
1 Parent(s): 7f8471e

PyTorch 1.6.0 compatibility updates

Browse files
Files changed (1) hide show
  1. 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