Pass options to TFE_ContextOptionsSetAsync

PiperOrigin-RevId: 190707017

diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py
index 6c9a147..8c1bb06 100644 (file)
--- a/tensorflow/python/eager/context.py
+++ b/tensorflow/python/eager/context.py
@@ -232,7 +232,7 @@ class Context(object):
             pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy(
                 opts, self._device_policy)
           if self._execution_mode == ASYNC:
-            pywrap_tensorflow.TFE_ContextOptionsSetAsync(True)
+            pywrap_tensorflow.TFE_ContextOptionsSetAsync(opts, True)
           self._context_handle = pywrap_tensorflow.TFE_NewContext(opts, status)
       finally:
         pywrap_tensorflow.TFE_DeleteContextOptions(opts)

diff --git a/tensorflow/python/eager/core_test.py b/tensorflow/python/eager/core_test.py
index 6ebf5b2..5f19f64 100644 (file)
--- a/tensorflow/python/eager/core_test.py
+++ b/tensorflow/python/eager/core_test.py
@@ -97,6 +97,14 @@ class TFETest(test_util.TensorFlowTestCase):
     self.assertTrue(has_cpu_device)
     del ctx
 
+  def testAsyncBasic(self):
+    ctx = context.Context(execution_mode=context.ASYNC)
+    has_cpu_device = False
+    for x in ctx.devices():
+      has_cpu_device = has_cpu_device or 'CPU' in x
+    self.assertTrue(has_cpu_device)
+    del ctx
+
   def testRunMetadata(self):
     context.enable_run_metadata()
     t = constant_op.constant(1.0)