Skip to content

GitLab

Commit bd0d5692 authored by Frank Chen's avatar Frank Chen Committed by A. Unique TensorFlower
Browse files

Consolidate DistributeOptions.auto_shard into DistributeOptions.auto_shard_policy. 

PiperOrigin-RevId: 275930249
parent aa3a6f23
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 2 deletions
official/nlp/bert/input_pipeline.py
View file @ bd0d5692
...@@ -51,7 +51,8 @@ def file_based_input_fn_builder(input_file, name_to_features): ...@@ -51,7 +51,8 @@ def file_based_input_fn_builder(input_file, name_to_features):
# same input file is sent to all workers. # same input file is sent to all workers.
if isinstance(input_file, str) or len(input_file) == 1: if isinstance(input_file, str) or len(input_file) == 1:
options = tf.data.Options() options = tf.data.Options()
options.experimental_distribute.auto_shard = False options.experimental_distribute.auto_shard_policy = (
tf.data.experimental.AutoShardPolicy.OFF)
d = d.with_options(options) d = d.with_options(options)
return d return d
......
official/nlp/xlnet/data_utils.py
View file @ bd0d5692
...@@ -119,7 +119,8 @@ def file_based_input_fn_builder(input_file, name_to_features, batch_size, ...@@ -119,7 +119,8 @@ def file_based_input_fn_builder(input_file, name_to_features, batch_size,
# same input file is sent to all workers. # same input file is sent to all workers.
if isinstance(input_file, str) or len(input_file) == 1: if isinstance(input_file, str) or len(input_file) == 1:
options = tf.data.Options() options = tf.data.Options()
options.experimental_distribute.auto_shard = False options.experimental_distribute.auto_shard_policy = (
tf.data.experimental.AutoShardPolicy.OFF)
d = d.with_options(options) d = d.with_options(options)
d = d.prefetch(tf.data.experimental.AUTOTUNE) d = d.prefetch(tf.data.experimental.AUTOTUNE)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment