Skip to content

GitLab

Commit 49b223b6 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 311048706
parent 36d27f01
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
official/recommendation/ncf_input_pipeline.py
View file @ 49b223b6
......@@ -45,7 +45,9 @@ def create_dataset_from_tf_record_files(input_file_pattern,
raise ValueError("Pre-batch ({}) size is not equal to batch "
"size ({})".format(pre_batch_size, batch_size))
files_dataset = files_dataset.shard(NUM_SHARDS, shard_index)
dataset = files_dataset.interleave(tf.data.TFRecordDataset)
dataset = files_dataset.interleave(
tf.data.TFRecordDataset,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
decode_fn = functools.partial(
data_pipeline.DatasetManager.deserialize,
batch_size=pre_batch_size,
......
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