Fix a bug

yzhangcs · yzhangcs · commit 4830efea201c · 2021-07-20T14:48:53.000+08:00
diff --git a/supar/models/dep.py b/supar/models/dep.py
@@ -558,13 +558,15 @@ def loss(self, s_arc, s_sib, s_rel, arcs, sibs, rels, mask, mbr=True, partial=Fa
                 original arc scores of shape ``[batch_size, seq_len, seq_len]`` if ``mbr=False``, or marginals otherwise.
         """
 
+        batch_size, seq_len = mask.shape
         if self.args.loss == 'crf':
             arc_dist = CRF2oDependency((s_arc, s_sib), mask, partial=partial)
             arc_loss = -arc_dist.log_prob((arcs, sibs)).sum() / mask.sum()
         elif self.args.loss == 'max-margin':
             s_arc = s_arc + torch.full_like(s_arc, 1).scatter_(-1, arcs.unsqueeze(-1), 0)
+            s_sib = s_sib + torch.full_like(s_sib, 1).masked_fill_(sibs.unsqueeze(-1).eq(sibs.new_tensor(range(seq_len))), 0)
             arc_dist = CRF2oDependency((s_arc, s_sib), mask, partial=partial)
-            arc_loss = (arc_dist.max - arc_dist.score(arcs)).sum() / mask.sum()
+            arc_loss = (arc_dist.max - arc_dist.score((arcs, sibs))).sum() / mask.sum()
         arc_probs = arc_dist.marginals if mbr else s_arc
         # -1 denotes un-annotated arcs
         if partial:
