yjh0410 1 ano atrás
pai
commit
e616dc7dc9
1 arquivos alterados com 1 adições e 1 exclusões
  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