Fix dummy batch when --max-tokens is small (fixes facebookresearch#347)

Summary: Pull Request resolved: facebookresearch#366 Differential Revision: D13058513 Pulled By: myleott fbshipit-source-id: a146d2cfb345d404775ed8d6b8e4a4ad4e7a33b4
moussaKam · Nov 14, 2018 · 3707c03 · 3707c03
1 parent 1f4e1e8
commit 3707c03
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fairseq/data/language_pair_dataset.py b/fairseq/data/language_pair_dataset.py
@@ -192,7 +192,7 @@ def get_dummy_batch(self, num_tokens, max_positions, src_len=128, tgt_len=128):
             max_positions,
             (self.max_source_positions, self.max_target_positions),
         )
-        bsz = num_tokens // max(src_len, tgt_len)
+        bsz = max(num_tokens // max(src_len, tgt_len), 1)
         return self.collater([
             {
                 'id': i,

diff --git a/fairseq/data/monolingual_dataset.py b/fairseq/data/monolingual_dataset.py
@@ -153,7 +153,7 @@ def get_dummy_batch(self, num_tokens, max_positions, tgt_len=128):
         """Return a dummy batch with a given number of tokens."""
         if isinstance(max_positions, float) or isinstance(max_positions, int):
             tgt_len = min(tgt_len, max_positions)
-        bsz = num_tokens // tgt_len
+        bsz = max(num_tokens // tgt_len, 1)
         target = self.vocab.dummy_sentence(tgt_len + 2)
         source, past_target, future_target = target[1:-1], target[2:], target[:-2]
         source, target = self._make_source_target(source, past_target, future_target)