glenn-jocher commited on
Commit
251aeaf
1 Parent(s): d2a1728

Update P5 + P6 model ensembling (#3082)

Browse files
Files changed (1) hide show
  1. models/experimental.py +3 -2
models/experimental.py CHANGED
@@ -130,7 +130,8 @@ def attempt_load(weights, map_location=None, inplace=True):
130
  if len(model) == 1:
131
  return model[-1] # return model
132
  else:
133
- print('Ensemble created with %s\n' % weights)
134
- for k in ['names', 'stride']:
135
  setattr(model, k, getattr(model[-1], k))
 
136
  return model # return ensemble
 
130
  if len(model) == 1:
131
  return model[-1] # return model
132
  else:
133
+ print(f'Ensemble created with {weights}\n')
134
+ for k in ['names']:
135
  setattr(model, k, getattr(model[-1], k))
136
+ model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride
137
  return model # return ensemble