Fix WongKinYiu#1101 (comment)

ijdoc · Apr 2, 2023 · 8875587 · 8875587
1 parent 55b90e1
commit 8875587
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/loss.py b/utils/loss.py
@@ -739,7 +739,7 @@ def build_targets(self, p, targets, imgs):
                 + 3.0 * pair_wise_iou_loss
             )
 
-            matching_matrix = torch.zeros_like(cost)
+            matching_matrix = torch.zeros_like(cost, device="cpu")
 
             for gt_idx in range(num_gt):
                 _, pos_idx = torch.topk(