OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-07 12:20:24 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Add check to tf.device when called with a function in eager mode.

Browse Source 
PiperOrigin-RevId: 173947845
This commit is contained in:
A. Unique TensorFlower 2017-10-30 14:16:18 -07:00 committed by TensorFlower Gardener
parent 3639aa7ff1
commit 1d6dae88ef
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/python/framework/ops.py
View File

@ -4339,11 +4339,18 @@ def device(device_name_or_function):
Returns:
A context manager that specifies the default device to use for newly
created ops.
Raises:
RuntimeError: If eager execution is enabled and a function is passed in.
"""
if context.in_graph_mode():
return get_default_graph().device(device_name_or_function)
else:
# TODO(agarwal): support device functions in EAGER mode.
if callable(device_name_or_function):
raise RuntimeError(
"tf.device does not support functions when eager execution "
"is enabled.")
return context.device(device_name_or_function)