Update P5 + P6 model ensembling (#3082)

ultralytics · May 8, 2021 · 251aeaf · 251aeaf
1 parent d2a1728
commit 251aeaf
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/models/experimental.py b/models/experimental.py
@@ -130,7 +130,8 @@ def attempt_load(weights, map_location=None, inplace=True):
     if len(model) == 1:
         return model[-1]  # return model
     else:
-        print('Ensemble created with %s\n' % weights)
-        for k in ['names', 'stride']:
+        print(f'Ensemble created with {weights}\n')
+        for k in ['names']:
             setattr(model, k, getattr(model[-1], k))
+        model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride  # max stride
         return model  # return ensemble