From 7f1d70d97f543d69a9f02cd6df0964f22f9278f3 Mon Sep 17 00:00:00 2001 From: Rohan Jain Date: Tue, 25 Sep 2018 20:16:49 -0700 Subject: 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 --- tensorflow/contrib/distribute/python/mirrored_strategy.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) (limited to 'tensorflow/contrib/distribute/python/mirrored_strategy.py') 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, -- cgit v1.2.3