From f70fded9c3bb5bdc1ef3b480dc3a48a9b9485e48 Mon Sep 17 00:00:00 2001 From: Pierre Pereira Date: Thu, 15 Jun 2023 15:03:42 +0200 Subject: [PATCH] fix: wrong grid device w.r.t. output --- yolox/models/yolo_head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yolox/models/yolo_head.py b/yolox/models/yolo_head.py index 3e51768ee..79e8ecc11 100644 --- a/yolox/models/yolo_head.py +++ b/yolox/models/yolo_head.py @@ -227,7 +227,7 @@ def get_output_and_grid(self, output, k, stride, dtype): output = output.permute(0, 1, 3, 4, 2).reshape( batch_size, hsize * wsize, -1 ) - grid = grid.view(1, -1, 2) + grid = grid.view(1, -1, 2).to(output.device) output[..., :2] = (output[..., :2] + grid) * stride output[..., 2:4] = torch.exp(output[..., 2:4]) * stride return output, grid