Internal change

PiperOrigin-RevId: 298986381

Internal change
PiperOrigin-RevId: 298986381
7a3d6c4c · Chen Chen · A. Unique TensorFlower · bb8a18c9 · 7a3d6c4c
Commit 7a3d6c4c authored Mar 04, 2020 by Chen Chen Committed by A. Unique TensorFlower Mar 04, 2020
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 9 deletions

official/nlp/bert/input_pipeline.py official/nlp/bert/input_pipeline.py +9 -9

No files found.
--- a/official/nlp/bert/input_pipeline.py
+++ b/official/nlp/bert/input_pipeline.py
@@ -87,7 +87,7 @@ def create_pretrain_dataset(input_patterns,
  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)
-
+  if is_training:
    dataset = dataset.repeat()

    # We set shuffle buffer to exactly match total number of
@@ -132,7 +132,7 @@ def create_pretrain_dataset(input_patterns,
  if is_training:
    dataset = dataset.shuffle(100)

-  dataset = dataset.batch(batch_size, drop_remainder=True)
+  dataset = dataset.batch(batch_size, drop_remainder=is_training)
  dataset = dataset.prefetch(1024)
  return dataset