diff options
| author |  Rohan Jain <rohanj@google.com> | 2018-09-25 20:16:49 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-25 20:22:00 -0700 |
| commit | 7f1d70d97f543d69a9f02cd6df0964f22f9278f3 (patch) | |
| tree | 29612b6cd40203beba4f2b9689eef27a1f8da8d7 /tensorflow/contrib/distribute/python/mirrored_strategy.py | |
| parent | 3f4b8c138165cc9deb0ed931c5a6bb3d8ab556f0 (diff) | |
Switching Distribution strategies to use MultiDeviceIterator. Currently only supported in Graph mode using initializable iterators. In a subsequent change, we'll add in support for Eager mode as well.
This removes prefetching_ops_v2 code.
PiperOrigin-RevId: 214546754
Diffstat (limited to 'tensorflow/contrib/distribute/python/mirrored_strategy.py')
| -rw-r--r-- | tensorflow/contrib/distribute/python/mirrored_strategy.py | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py index 0c6805d682..945f450387 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -480,8 +480,10 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): self._prefetch_on_device) else: return values.PerDeviceDataset( - self._call_dataset_fn(dataset_fn), self._devices, - self._prefetch_on_device) + self._call_dataset_fn(dataset_fn), + self._devices, + self._prefetch_on_device, + source_device=device_util.resolve("/device:CPU:0")) # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. def _run_steps_on_dataset(self, fn, iterator, iterations, |