Browse Source

DetectMultiBackend() return `device` update (#6958)

Fixes ONNX validation that returns outputs on CPU.
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
52c1399fdc
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
  1. +2
    -1
      models/common.py

+ 2
- 1
models/common.py View File

@@ -458,7 +458,8 @@ class DetectMultiBackend(nn.Module):
y = (y.astype(np.float32) - zero_point) * scale # re-scale
y[..., :4] *= [w, h, w, h] # xywh normalized to pixels

y = torch.tensor(y) if isinstance(y, np.ndarray) else y
if isinstance(y, np.ndarray):
y = torch.tensor(y, device=self.device)
return (y, []) if val else y

def warmup(self, imgsz=(1, 3, 640, 640)):

Loading…
Cancel
Save