Merge 52a40ca into 2d27cad

harvardnlp · Mar 22, 2020 · d2938ad · d2938ad
2 parents 2d27cad + 52a40ca
commit d2938ad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_struct/deptree.py b/torch_struct/deptree.py
@@ -13,7 +13,7 @@ def _convert(logits):
 
     N = logits.size(1)
     new_logits[:, 0, 1:] = logits[:, torch.arange(N), torch.arange(N)]
-    new_logits[:, torch.arange(1, N), torch.arange(1, N)] = -1e9
+    new_logits[:, torch.arange(1, N + 1), torch.arange(1, N + 1)] = -1e9
     return new_logits