Enhanced dependency parsing

combo/models/graph_parser.py

@@ -164,9 +164,9 @@ class GraphDependencyRelationModel(base.Predictor):
               heads_true: torch.Tensor,
               mask: torch.BoolTensor,
               sample_weights: torch.Tensor) -> torch.Tensor:
-
-        true = true[true.long() > 0]
-        pred = pred[heads_true.long() == 1]
+        correct_heads_mask = heads_true.long() == 1
+        true = true[correct_heads_mask]
+        pred = pred[correct_heads_mask]
         loss = F.cross_entropy(pred, true.long())
         return loss.sum() / pred.size(0)