[Enhance]: Avoid crash in empty gt training of GFL (#4631)

* avoid crash in empty gt training of GFL * real fix of the bug
4 years ago · 2a24412846
parent 283f82098b
commit 2a24412846
1 changed files with 1 additions and 1 deletions
--- a/mmdet/models/dense_heads/gfl_head.py
+++ b/mmdet/models/dense_heads/gfl_head.py
@ -284,7 +284,7 @@ class GFLHead(AnchorHead):
        else:
            loss_bbox = bbox_pred.sum() * 0
            loss_dfl = bbox_pred.sum() * 0
-            weight_targets = torch.tensor(0).cuda()
+            weight_targets = bbox_pred.new_tensor(0)

        # cls (qfl) loss
        loss_cls = self.loss_cls(