fix ddp unused param bug

naive0409 · Nov 9, 2022 · 5738ccf · 5738ccf
1 parent 1563345
commit 5738ccf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/yolov6/models/loss.py b/yolov6/models/loss.py
@@ -248,11 +248,11 @@ def forward(self, pred_dist, pred_bboxes, anchor_points,
                 else:
                     loss_dfl = loss_dfl.sum() / target_scores_sum
             else:
-                loss_dfl = torch.tensor(0.).to(pred_dist.device)
+                loss_dfl = pred_dist.sum() * 0.
 
         else:
-            loss_iou = torch.tensor(0.).to(pred_dist.device)
-            loss_dfl = torch.tensor(0.).to(pred_dist.device)
+            loss_iou = pred_dist.sum() * 0.
+            loss_dfl = pred_dist.sum() * 0.
 
         return loss_iou, loss_dfl