diff options
| author |  Rohan Jain <rohanj@google.com> | 2018-09-27 20:59:37 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-27 21:02:52 -0700 |
| commit | 97cba0b88cb3ce6a3f3cc66a8c4fd414bd3ac1a8 (patch) | |
| tree | f7ce9c26ebb93cbef7aacc3f725ba5a304470407 /tensorflow/contrib/distribute/python/mirrored_strategy.py | |
| parent | 370d385c3029a7972ba201c8303942b30f09521c (diff) | |
Allowing source_device to be set to /cpu:0 for multi device iterator in distribution strategies. That is always the appropriate option.
In the existing code, we would set it to a partially specified "worker" name that was ambiguous and end up on the GPU.
PiperOrigin-RevId: 214882658
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 945f450387..504f45a695 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -482,8 +482,7 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): return values.PerDeviceDataset( self._call_dataset_fn(dataset_fn), self._devices, - self._prefetch_on_device, - source_device=device_util.resolve("/device:CPU:0")) + self._prefetch_on_device) # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. def _run_steps_on_dataset(self, fn, iterator, iterations, |