fix batch size (#462)

pytorch · Jan 30, 2019 · 1c2ae32 · 1c2ae32
1 parent 1f64e82
commit 1c2ae32
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/torchtext/data/iterator.py b/torchtext/data/iterator.py
@@ -220,10 +220,15 @@ def __iter__(self):
             for i in range(0, len(self) * self.bptt_len, self.bptt_len):
                 self.iterations += 1
                 seq_len = min(self.bptt_len, len(data) - i - 1)
+                batch_text = data[i:i + seq_len]
+                batch_target = data[i + 1:i + 1 + seq_len]
+                if TEXT.batch_first:
+                    batch_text = batch_text.t().contiguous()
+                    batch_target = batch_target.t().contiguous()
                 yield Batch.fromvars(
                     dataset, self.batch_size,
-                    text=data[i:i + seq_len],
-                    target=data[i + 1:i + 1 + seq_len])
+                    text=batch_text,
+                    target=batch_target)
             if not self.repeat:
                 return