.

harvardnlp · Sep 3, 2019 · eea389c · eea389c
1 parent 970c3e1
commit eea389c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/torch_struct/data/data.py b/torch_struct/data/data.py
@@ -57,15 +57,15 @@ def SubTokenizedField(tokenizer):
     FIELD.is_target = False
     return FIELD
 
-def TokenBucket(train):
+def TokenBucket(train, batch_size):
     def batch_size_fn(x, _, size):
         return size + max(len(x.word[0]), 5)
     return torchtext.data.BucketIterator(train,
         train=True,
         sort=False,
         sort_within_batch=True,
         shuffle=True,
-        batch_size=1500,
+        batch_size=batch_size,
         sort_key=lambda x: len(x.word[0]),
         repeat=True,
         batch_size_fn=batch_size_fn,

diff --git a/torch_struct/deptree.py b/torch_struct/deptree.py
@@ -141,8 +141,8 @@ def deptree_fromseq(sequence, lengths=None):
     if lengths is None:
         lengths = torch.LongTensor([N] * batch)
     labels = torch.zeros(batch, N + 1, N + 1).long()
-    for n in range(1, N):
-        labels[torch.arange(batch), sequence[:, n], n] = 1
+    for n in range(1, N+1):
+        labels[torch.arange(batch), sequence[:, n-1], n] = 1
     for b in range(batch):
         labels[b, lengths[b]+1:, :] = 0
         labels[b, :, lengths[b]+1:] = 0