diff options
| author |  Akshay Modi <nareshmodi@google.com> | 2018-03-27 18:18:33 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-03-27 18:21:13 -0700 |
| commit | 9d0b5ba5cfdf0f1c35a228abac9c84c11bb7d695 (patch) | |
| tree | 3e5f864b7001f9c68a2c1a5cd86406bdf6bd974e | |
| parent | 94b2d2db576a6cce878aee92d6b1f90ded4278b4 (diff) | |
Pass options to TFE_ContextOptionsSetAsync
PiperOrigin-RevId: 190707017
| -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) |