From 0a708f8da4afc5099cae1493dfb60a5680dadf2f Mon Sep 17 00:00:00 2001
From: Jianwei Xie <xiejw@google.com>
Date: Sat, 22 Dec 2018 21:40:58 -0800
Subject: [PATCH] Re-prefetch the dataset for V1Adapater

PiperOrigin-RevId: 226650870
---
 tensorflow/python/distribute/values.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
index ef1d2a94992..1f5077a75ae 100644
--- a/tensorflow/python/distribute/values.py
+++ b/tensorflow/python/distribute/values.py
@@ -1959,6 +1959,9 @@ def _split_dataset_batch(dataset, split_batch_by):
   prefetch_buffer = None
   if isinstance(dataset, dataset_ops.PrefetchDataset):
     prefetch_buffer = dataset._buffer_size
+  elif (isinstance(dataset, dataset_ops.DatasetV1Adapter)
+        and isinstance(dataset._dataset, dataset_ops.PrefetchDataset)):
+    prefetch_buffer = dataset._dataset._buffer_size
   # pylint: enable=protected-access
 
   if tensor_util.is_tensor(batch_size):