glenn-jocher commited on
Commit
199c9c7
1 Parent(s): 05a955a

Detection() device bug fix (#1455)

Browse files
Files changed (1) hide show
  1. models/common.py +2 -1
models/common.py CHANGED
@@ -183,7 +183,8 @@ class Detections:
183
  self.names = names # class names
184
  self.xyxy = pred # xyxy pixels
185
  self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
186
- gn = [torch.Tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.]) for im in imgs] # normalization gains
 
187
  self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
188
  self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
189
  self.n = len(self.pred)
 
183
  self.names = names # class names
184
  self.xyxy = pred # xyxy pixels
185
  self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
186
+ d = pred[0].device # device
187
+ gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
188
  self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
189
  self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
190
  self.n = len(self.pred)