fix(model): yolo_head.py not support fp16 when fallback to cpu mode (#…

…688)
Megvii-BaseDetection · Sep 13, 2021 · be02441 · be02441
1 parent 6819425
commit be02441
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/yolox/models/yolo_head.py b/yolox/models/yolo_head.py
@@ -489,7 +489,7 @@ def get_assignments(
         with torch.cuda.amp.autocast(enabled=False):
             cls_preds_ = (
                 cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
-                * obj_preds_.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
+                * obj_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
             )
             pair_wise_cls_loss = F.binary_cross_entropy(
                 cls_preds_.sqrt_(), gt_cls_per_image, reduction="none"