aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/distribute/python/mirrored_strategy.py
diff options
context:
space:
mode:
authorGravatar Rohan Jain <rohanj@google.com>2018-09-25 20:16:49 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-09-25 20:22:00 -0700
commit7f1d70d97f543d69a9f02cd6df0964f22f9278f3 (patch)
tree29612b6cd40203beba4f2b9689eef27a1f8da8d7 /tensorflow/contrib/distribute/python/mirrored_strategy.py
parent3f4b8c138165cc9deb0ed931c5a6bb3d8ab556f0 (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.py6
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,