Fix NHNet input pipeline.

PiperOrigin-RevId: 310251449

Fix NHNet input pipeline.
PiperOrigin-RevId: 310251449
92903d58 · Hongkun Yu · A. Unique TensorFlower · 3097fd2a · 92903d58
Commit 92903d58 authored May 06, 2020 by Hongkun Yu Committed by A. Unique TensorFlower May 06, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/nlp/nhnet/input_pipeline.py official/nlp/nhnet/input_pipeline.py +1 -1

No files found.
--- a/official/nlp/nhnet/input_pipeline.py
+++ b/official/nlp/nhnet/input_pipeline.py
@@ -208,7 +208,7 @@ def create_dataset(file_paths,
    else:
      dataset = process_singledoc_transformer_dataset(dataset, batch_size,
                                                      params)
-  dataset = dataset.prefetch(1024)
+  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  return dataset