aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/distribute/python/mirrored_strategy.py
diff options
context:
space:
mode:
authorGravatar Priya Gupta <priyag@google.com>2018-04-18 12:03:32 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-04-18 12:06:14 -0700
commit03d18ae232c3cff4c56d1efec7bf29f9b16c4f68 (patch)
treeb722f735de3f6ab2d3e8ef945baab452bd0b70ab /tensorflow/contrib/distribute/python/mirrored_strategy.py
parent60444df318439654324ff797d66734c9920e48a2 (diff)
Add support for initializable iterator in distribution strategies. Use that in estimator.
PiperOrigin-RevId: 193394603
Diffstat (limited to 'tensorflow/contrib/distribute/python/mirrored_strategy.py')
-rw-r--r--tensorflow/contrib/distribute/python/mirrored_strategy.py3
1 files changed, 1 insertions, 2 deletions
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py
index eb0edb3a11..d5e22e8100 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py
@@ -141,9 +141,8 @@ class MirroredStrategy(distribute_lib.DistributionStrategy):
return result
def distribute_dataset(self, dataset):
- per_device_dataset = values.PerDeviceDataset(
+ return values.PerDeviceDataset(
dataset, self._devices, self._prefetch_on_device)
- return per_device_dataset.make_one_shot_iterator()
def _broadcast(self, tensor, destinations):
# TODO(josh11b): In eager mode, use one thread per device, or async mode.