Skip to content

Commit 3f8604b

Browse files
fix(mmdet/base_dense_head): 🐛 correct slicing issue in score softmax for single-class models (open-mmlab#2827)
The softmax function previously sliced out the background class, which caused an IndexError in single-class model export. This update removes the slicing of background class scores, fixing the issue when exporting RetinaNet with single-class configurations. - Updated `base_dense_head__predict_by_feat` to prevent IndexError when handling single-class outputs. - This change is particularly relevant for RetinaNet model exports.
1 parent bc75c9d commit 3f8604b

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ def base_dense_head__predict_by_feat(
113113
if self.use_sigmoid_cls:
114114
scores = scores.sigmoid()
115115
else:
116-
scores = scores.softmax(-1)[:, :, :-1]
116+
scores = scores.softmax(-1)
117117
if with_score_factors:
118118
score_factors = score_factors.permute(0, 2, 3,
119119
1).reshape(batch_size,

0 commit comments

Comments
 (0)