Detection() device bug fix (#1455)

ultralytics · Nov 19, 2020 · 199c9c7 · 199c9c7
1 parent 05a955a
commit 199c9c7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/models/common.py b/models/common.py
@@ -183,7 +183,8 @@ def __init__(self, imgs, pred, names=None):
         self.names = names  # class names
         self.xyxy = pred  # xyxy pixels
         self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
-        gn = [torch.Tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.]) for im in imgs]  # normalization gains
+        d = pred[0].device  # device
+        gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs]  # normalizations
         self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
         self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
         self.n = len(self.pred)