Merge pull request meituan#605 from meituan-gengyifei/main

Fix the bug of unused parameters in DDP mode.
naive0409 · Nov 9, 2022 · d6d89b6 · d6d89b6
2 parents 1563345 + 39996fd
commit d6d89b6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 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
 

diff --git a/yolov6/models/loss_distill.py b/yolov6/models/loss_distill.py
@@ -324,14 +324,14 @@ def forward(self, pred_dist, pred_bboxes, t_pred_dist, t_pred_bboxes, temperatur
                     loss_dfl = loss_dfl.sum() / target_scores_sum
                     d_loss_dfl = d_loss_dfl.sum() / target_scores_sum
             else:
-                loss_dfl = torch.tensor(0.).to(pred_dist.device)
-                d_loss_dfl = torch.tensor(0.).to(pred_dist.device)
+                loss_dfl = pred_dist.sum() * 0.
+                d_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)
-            d_loss_dfl = torch.tensor(0.).to(pred_dist.device)
+            loss_iou = pred_dist.sum() * 0.
+            loss_dfl = pred_dist.sum() * 0.
+            d_loss_dfl = pred_dist.sum() * 0.
 
         return loss_iou, loss_dfl, d_loss_dfl