Disabling MultiDeviceIterator in Eager mode. Support is coming soon.

PiperOrigin-RevId: 214296771

1 files changed, 15 insertions, 1 deletions

diff --git a/tensorflow/python/data/ops/multi_device_iterator_ops.py b/tensorflow/python/data/ops/multi_device_iterator_ops.py
index 84e8abbd83..c914a43956 100644
--- a/tensorflow/python/data/ops/multi_device_iterator_ops.py
+++ b/tensorflow/python/data/ops/multi_device_iterator_ops.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.data.util import nest
 from tensorflow.python.data.util import sparse
+from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
@@ -129,7 +130,13 @@ class _PerDeviceGenerator(dataset_ops.Dataset):
 
 
 class MultiDeviceIterator(object):
-  """An iterator over multiple devices."""
+  """An iterator over multiple devices.
+
+  @compatibility(eager)
+  MultiDeviceIterator isn't currently supported in Eager mode but support is
+  coming soon.
+  @end_compatibility
+  """
 
   def __init__(self,
                dataset,
@@ -146,7 +153,14 @@ class MultiDeviceIterator(object):
       prefetch_buffer_size: if > 1, then we setup a buffer on each device
         to prefetch into.
       source_device: The host device to place the `dataset` on.
+
+    Raises:
+      RuntimeError: If run in Eager mode.
     """
+    if context.executing_eagerly():
+      # TODO(rohanj): Fix this. Tracking bug: b/116467184
+      raise RuntimeError("MultiDeviceIterator is not currently supported in "
+                         "Eager mode.")
     self._dataset = dataset
     self._devices = devices
     self._source_device = source_device