Fix torch long to float tensor on HUB macOS (ultralytics#8067)

tdhooghe · Jun 10, 2022 · ffe6b25 · ffe6b25
1 parent c1c6dca
commit ffe6b25
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/utils/loss.py b/utils/loss.py
@@ -195,8 +195,8 @@ def build_targets(self, p, targets):
             device=self.device).float() * g  # offsets
 
         for i in range(self.nl):
-            anchors = self.anchors[i]
-            gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain
+            anchors, shape = self.anchors[i], p[i].shape
+            gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain
 
             # Match targets to anchors
             t = targets * gain  # shape(3,n,7)
@@ -226,7 +226,7 @@ def build_targets(self, p, targets):
             gi, gj = gij.T  # grid indices
 
             # Append
-            indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
+            indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid
             tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
             anch.append(anchors[a])  # anchors
             tcls.append(c)  # class