Skip to content

Commit

Permalink
fix(mmdet/base_dense_head): 🐛 correct slicing issue in score softmax …
Browse files Browse the repository at this point in the history
…for single-class models (#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.
  • Loading branch information
rchuzh99 authored Sep 30, 2024
1 parent bc75c9d commit 3f8604b
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ def base_dense_head__predict_by_feat(
if self.use_sigmoid_cls:
scores = scores.sigmoid()
else:
scores = scores.softmax(-1)[:, :, :-1]
scores = scores.softmax(-1)
if with_score_factors:
score_factors = score_factors.permute(0, 2, 3,
1).reshape(batch_size,
Expand Down

0 comments on commit 3f8604b

Please sign in to comment.