Skip to content

GitLab

Commit b4af1c94 authored by zehaos's avatar zehaos
Browse files

Fix index error when class_agnostic is true.

parent c95c6373
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 1 deletion
mmdet/models/mask_heads/fcn_mask_head.py
View file @ b4af1c94
...@@ -97,6 +97,10 @@ class FCNMaskHead(nn.Module): ...@@ -97,6 +97,10 @@ class FCNMaskHead(nn.Module):
def loss(self, mask_pred, mask_targets, labels): def loss(self, mask_pred, mask_targets, labels):
loss = dict() loss = dict()
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_mask = mask_cross_entropy(mask_pred, mask_targets, labels)
loss['loss_mask'] = loss_mask loss['loss_mask'] = loss_mask
return loss return loss
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment