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