fix: fix bugs of num_pred in DeformableDETRHead

yumion · Sep 20, 2022 · 2edd0c9 · 2edd0c9
1 parent d0b5aeb
commit 2edd0c9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mmdet/models/dense_heads/deformable_detr_head.py b/mmdet/models/dense_heads/deformable_detr_head.py
@@ -62,8 +62,8 @@ def _init_layers(self) -> None:
 
         # last reg_branch is used to generate proposal from
         # encode feature map when as_two_stage is True.
-        num_pred = (self.transformer.decoder.num_layers + 1) if \
-            self.as_two_stage else self.num_decoder_layers #TODO
+        num_pred = (self.num_decoder_layers + 1) if \
+            self.as_two_stage else self.num_decoder_layers
 
         if self.with_box_refine:
             self.cls_branches = nn.ModuleList(