yjh0410 1 年間 前
コミット
e616dc7dc9
1 ファイル変更1 行追加1 行削除
  1. 1 1
      odlab/models/detectors/fcos/criterion.py

+ 1 - 1
odlab/models/detectors/fcos/criterion.py

@@ -119,7 +119,7 @@ class SetCriterion(nn.Module):
         loss_box = 1.0 - ious
 
         if box_weight is not None:
-            loss_box = loss_box.unsqueeze(1) * box_weight
+            loss_box = loss_box.squeeze(-1) * box_weight
 
         return loss_box.sum() / num_boxes