diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2018-04-18 16:35:44 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-18 16:38:07 -0700 |
| commit | fddfa9f8dcd1a922ade5362c0538ca39e99472a7 (patch) | |
| tree | 7320668c20041c85bcce4e0c6614f7270dce44c6 /tensorflow/contrib/distribute/python/mirrored_strategy.py | |
| parent | e9d47fbff0d644a75c6f3dcdcb852685ef515b64 (diff) | |
Change distribution.distribute_dataset to accept an input_fn instead of a dataset.
PiperOrigin-RevId: 193437651
Diffstat (limited to 'tensorflow/contrib/distribute/python/mirrored_strategy.py')
| -rw-r--r-- | tensorflow/contrib/distribute/python/mirrored_strategy.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py index d5e22e8100..6efd578a77 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -140,9 +140,10 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): g.add_to_collections(collections, result) return result - def distribute_dataset(self, dataset): + def distribute_dataset(self, dataset_fn): return values.PerDeviceDataset( - dataset, self._devices, self._prefetch_on_device) + self._call_dataset_fn(dataset_fn), self._devices, + self._prefetch_on_device) def _broadcast(self, tensor, destinations): # TODO(josh11b): In eager mode, use one thread per device, or async mode. |