Merge pull request iscyy#51 from ThibaultCastells/fix_cast_error_v2

Fix for "RuntimeError: result type Float can't be cast to the desired output type long int"
lzu-cvpr · Sep 7, 2022 · 2ae3bea · 2ae3bea
2 parents 9cddc40 + 7582f9f
commit 2ae3bea
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/utils/loss.py b/utils/loss.py
@@ -186,8 +186,8 @@ def build_targets(self, p, targets):
                             ], device=targets.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
@@ -219,7 +219,7 @@ def build_targets(self, p, targets):
 
             # Append
             a = t[:, 6].long()  # anchor indices
-            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 indices
             tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
             anch.append(anchors[a])  # anchors
             tcls.append(c)  # class