Update loss for FP16 tobj (#7088)
This commit is contained in:
1 parent
6f128031d0
commit
a2d617ece9
1 file changed
+1
-1
+1
-1
@@ -125,7 +125,7 @@ class ComputeLoss:
|
|||||||
# Losses
|
# Losses
|
||||||
for i, pi in enumerate(p): # layer index, layer predictions
|
for i, pi in enumerate(p): # layer index, layer predictions
|
||||||
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
||||||
tobj = torch.zeros(pi.shape[:4], device=self.device) # target obj
|
tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
|
||||||
|
|
||||||
n = b.shape[0] # number of targets
|
n = b.shape[0] # number of targets
|
||||||
if n:
|
if n:
|
||||||
|
|||||||
Reference in new issue
Block a user