Enable usage of batch_by_size for other code paths than just train (#17)

pytorch-tpu · Nov 7, 2019 · aa2c3b3 · aa2c3b3
1 parent 92f19a2
commit aa2c3b3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/fairseq/tasks/fairseq_task.py b/fairseq/tasks/fairseq_task.py
@@ -138,9 +138,17 @@ def get_batch_iterator(
             )
 
         # create mini-batches with given size constraints
-        batch_sampler = data_utils.batch_by_size_tpu(
-            indices, dataset.num_tokens, self.args.input_shapes
-        )
+        if getattr(self.args, 'use_gpu', True):
+            batch_sampler = data_utils.batch_by_size(
+                indices, dataset.num_tokens, max_tokens=max_tokens,
+                max_sentences=max_sentences,
+                required_batch_size_multiple=required_batch_size_multiple,
+            )
+        else:
+            batch_sampler = data_utils.batch_by_size_tpu(
+                indices, dataset.num_tokens,
+                getattr(self.args, 'input_shapes', None)
+            )
 
         # return a reusable, sharded iterator
         return iterators.EpochBatchIterator(

diff --git a/fairseq/tasks/translation.py b/fairseq/tasks/translation.py
@@ -188,7 +188,7 @@ def load_dataset(self, split, epoch=0, combine=False, **kwargs):
             left_pad_target=self.args.left_pad_target,
             max_source_positions=self.args.max_source_positions,
             max_target_positions=self.args.max_target_positions,
-            input_shapes=self.args.input_shapes,
+            input_shapes=getattr(self.args, 'input_shapes', None),
         )
 
     def build_dataset_for_inference(self, src_tokens, src_lengths):