Update dynamic_mask_head.py

stevewongv · Jul 18, 2022 · aa6ddc7 · aa6ddc7
1 parent a1b058e
commit aa6ddc7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/adet/modeling/ssis/dynamic_mask_head.py b/adet/modeling/ssis/dynamic_mask_head.py
@@ -290,6 +290,8 @@ def convert_locations(relative_coords,H,W):
             asso_mask_logits = asso_mask_logits.reshape(-1, 2, H, W)
             boundary_logits = mask_logits[:,1:2] 
             asso_boundary_logits = asso_mask_logits[:,1:2] 
+            mask_logits = mask_logits[:,0:1] 
+            asso_mask_logits = asso_mask_logits[:,0:1]
         else:
             mask_logits = mask_logits.reshape(-1, 1, H, W)
             asso_mask_logits = asso_mask_logits.reshape(-1, 1, H, W)