.

harvardnlp · Oct 21, 2019 · 516d958 · 516d958
1 parent b770f21
commit 516d958
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/torch_struct/cky_crf.py b/torch_struct/cky_crf.py
@@ -25,10 +25,8 @@ def _dp(self, scores, lengths=None, force_grad=False):
         for w in range(1, N):
             Y = beta[A][:, :, : N - w, :w]
             Z = beta[B][:, :, w:, N - w :]
-            f = torch.arange(N - w)
-            X = reduced_scores[:, :, f, f + w]
-
-            beta[A][:, :, : N - w, w] = semiring.times(semiring.dot(Y, Z), X)
+            score = torch.diagonal(reduced_scores, w, 2, 3)
+            beta[A][:, :, : N - w, w] = semiring.times(semiring.dot(Y, Z), score)
             beta[B][:, :, w:N, N - w - 1] = beta[A][:, :, : N - w, w]
 
         final = beta[A][:, :, 0]