Update train_embedding.py (#28)

titu1994 · Aug 27, 2021 · 5885441 · 5885441
1 parent f35d964
commit 5885441
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/train_embedding.py b/examples/train_embedding.py
@@ -29,11 +29,11 @@ def transform(_):
 # Note: Train and test are drawn from same distribution for demonstration purposes.
 # We should get near identical scores on both of them.
 ds_train = ds_full.take(train_size)
-ds_train = ds_train.map(transform)
+ds_train = ds_train.map(transform,num_parallel_calls=tf.data.experimental.AUTOTUNE)
 ds_train = ds_train.batch(BATCH_SIZE)
 
 ds_test = ds_full.skip(train_size)
-ds_test = ds_test.map(transform)
+ds_test = ds_test.map(transform,num_parallel_calls=tf.data.experimental.AUTOTUNE)
 ds_test = ds_test.batch(BATCH_SIZE)