diff options
| -rw-r--r-- | tensorflow/python/eager/context.py | 2 | ||||
| -rw-r--r-- | tensorflow/python/eager/core_test.py | 8 |
2 files changed, 9 insertions, 1 deletions
diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py index 6c9a14730c..8c1bb06bc3 100644 --- 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 6ebf5b2481..5f19f64846 100644 --- 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) |