commit 6fdaa569ac5a45c610dbfb59720c4441581bdd2e
author DEKHTIARJonathan <contact@jonathandekhtiar.eu> Tue Apr 07 16:43:47 2020 -0700
committer DEKHTIARJonathan <contact@jonathandekhtiar.eu> Tue Apr 07 16:43:47 2020 -0700
tree b48482f0144c056140306395e9c881ff89f1bb19
parent 6dcbc61bfb05c87aa9b184a2feaf189e39fbe58d

Change ops.colocate_with to ops.device in Iterator - 2

tensorflow/python/data/ops/iterator_ops.py

diff --git a/tensorflow/python/data/ops/iterator_ops.py b/tensorflow/python/data/ops/iterator_ops.py
index bd904d3..d69f33a 100644
--- a/tensorflow/python/data/ops/iterator_ops.py
+++ b/tensorflow/python/data/ops/iterator_ops.py
@@ -368,7 +368,7 @@
                           "dataset with output shapes %r." %
                           (self.output_shapes, dataset_output_shapes))
 
-    with ops.colocate_with(dataset._variant_tensor):
+    with ops.device(dataset._variant_tensor.device):
       return gen_dataset_ops.make_iterator(
           dataset._variant_tensor, self._iterator_resource, name=name)  # pylint: disable=protected-access