Fix index error when class_agnostic is true.

b4af1c94 · zehaos · c95c6373 · b4af1c94
Commit b4af1c94 authored Dec 12, 2018 by zehaos
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

mmdet/models/mask_heads/fcn_mask_head.py mmdet/models/mask_heads/fcn_mask_head.py +5 -1

No files found.
--- a/mmdet/models/mask_heads/fcn_mask_head.py
+++ b/mmdet/models/mask_heads/fcn_mask_head.py
@@ -97,7 +97,11 @@ class FCNMaskHead(nn.Module):

    def loss(self, mask_pred, mask_targets, labels):
        loss = dict()
-        loss_mask = mask_cross_entropy(mask_pred, mask_targets, labels)
+        if self.class_agnostic:
+            loss_mask = mask_cross_entropy(mask_pred, mask_targets,
+                                           torch.zeros_like(labels))
+        else:
+            loss_mask = mask_cross_entropy(mask_pred, mask_targets, labels)
        loss['loss_mask'] = loss_mask
        return loss