Update cky.py

srush · web-flow · commit ede76f8d3628 · 2020-04-11T14:27:24.000-04:00
diff --git a/torch_struct/cky.py b/torch_struct/cky.py
@@ -23,7 +23,7 @@ def _dp(self, scores, lengths=None, force_grad=False, cache=True):
             semiring.convert(roots).requires_grad_(True),
         )
         if lengths is None:
-            lengths = torch.LongTensor([N] * batch)
+            lengths = torch.LongTensor([N] * batch).to(terms.device)
 
         # Charts
         beta = [

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ def _dp(self, scores, lengths=None, force_grad=False, cache=True):`
`23`	`23`	`semiring.convert(roots).requires_grad_(True),`
`24`	`24`	`)`
`25`	`25`	`if lengths is None:`
`26`		`- lengths = torch.LongTensor([N] * batch)`
	`26`	`+ lengths = torch.LongTensor([N] * batch).to(terms.device)`
`27`	`27`
`28`	`28`	`# Charts`
`29`	`29`	`beta = [`