glenn-jocher
commited on
Commit
•
7cef03d
1
Parent(s):
a80dd66
Fix torch `long` to `float` tensor on HUB macOS (#8067)
Browse files- utils/loss.py +3 -3
utils/loss.py
CHANGED
@@ -195,8 +195,8 @@ class ComputeLoss:
|
|
195 |
device=self.device).float() * g # offsets
|
196 |
|
197 |
for i in range(self.nl):
|
198 |
-
anchors = self.anchors[i]
|
199 |
-
gain[2:6] = torch.tensor(
|
200 |
|
201 |
# Match targets to anchors
|
202 |
t = targets * gain # shape(3,n,7)
|
@@ -226,7 +226,7 @@ class ComputeLoss:
|
|
226 |
gi, gj = gij.T # grid indices
|
227 |
|
228 |
# Append
|
229 |
-
indices.append((b, a, gj.clamp_(0,
|
230 |
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
|
231 |
anch.append(anchors[a]) # anchors
|
232 |
tcls.append(c) # class
|
|
|
195 |
device=self.device).float() * g # offsets
|
196 |
|
197 |
for i in range(self.nl):
|
198 |
+
anchors, shape = self.anchors[i], p[i].shape
|
199 |
+
gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain
|
200 |
|
201 |
# Match targets to anchors
|
202 |
t = targets * gain # shape(3,n,7)
|
|
|
226 |
gi, gj = gij.T # grid indices
|
227 |
|
228 |
# Append
|
229 |
+
indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid
|
230 |
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
|
231 |
anch.append(anchors[a]) # anchors
|
232 |
tcls.append(c) # class
|