Update yolo.py

ultralytics · Oct 2, 2021 · d40cad5 · d40cad5
1 parent 6bdd8b7
commit d40cad5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/models/yolo.py b/models/yolo.py
@@ -172,13 +172,13 @@ def _clip_augmented(self, y):
         nl = self.model[-1].nl  # number of detection layers (P3-P5)
         n = sum(4 ** x for x in range(nl))
 
-        # small
+        # large
         i = y[0].shape[1] // n
-        y[0] = y[0][:, i:]
+        y[0] = y[0][:, :-i]
 
-        # large
+        # small
         i = (y[-1].shape[1] // n) * 4 ** (nl - 1)
-        y[-1] = y[-1][:, :-i]
+        y[-1] = y[-1][:, i:]
 
         return y