diff options
author | Priya Gupta <priyag@google.com> | 2018-04-18 12:03:32 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-18 12:06:14 -0700 |
commit | 03d18ae232c3cff4c56d1efec7bf29f9b16c4f68 (patch) | |
tree | b722f735de3f6ab2d3e8ef945baab452bd0b70ab /tensorflow/contrib/distribute/python/mirrored_strategy.py | |
parent | 60444df318439654324ff797d66734c9920e48a2 (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.py | 3 |
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. |