diff --git a/utils/loss.py b/utils/loss.py index b49cc7f66e66..a06330e034bc 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -125,7 +125,7 @@ def __call__(self, p, targets): # predictions, targets # Losses for i, pi in enumerate(p): # layer index, layer predictions 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 if n: