Update linearchain.py

harvardnlp · Apr 11, 2020 · 5e8d662 · 5e8d662
1 parent fba3d5b
commit 5e8d662
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_struct/linearchain.py b/torch_struct/linearchain.py
@@ -104,7 +104,7 @@ def to_parts(sequence, extra, lengths=None):
         batch, N = sequence.shape
         labels = torch.zeros(batch, N - 1, C, C).long()
         if lengths is None:
-            lengths = torch.LongTensor([N] * batch).to(edge.device)
+            lengths = torch.LongTensor([N] * batch)
         for n in range(1, N):
             labels[torch.arange(batch), n - 1, sequence[:, n], sequence[:, n - 1]] = 1
         for b in range(batch):