From 199c9c787427ece5723d5309e1c7c524a99bc59d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 19 Nov 2020 13:45:36 +0100 Subject: [PATCH] Detection() device bug fix (#1455) --- models/common.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/common.py b/models/common.py index 6c3eea6fd475..78469695b9a2 100644 --- 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)