yjh0410 1 år sedan
förälder
incheckning
e616dc7dc9
1 ändrade filer med 1 tillägg och 1 borttagningar
  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