Wrapper for `Graph.device()` using the default graph. See `tf.Graph.device` for more details. Args: device_name_or_function: The device name or function to use in the context. Returns: A context manager that specifies the default device to use for newly created ops. Raises:
(device_name_or_function)
| 5173 | |
| 5174 | @tf_export(v1=["device"]) |
| 5175 | def device(device_name_or_function): |
| 5176 | """Wrapper for `Graph.device()` using the default graph. |
| 5177 | |
| 5178 | See `tf.Graph.device` for more details. |
| 5179 | |
| 5180 | Args: |
| 5181 | device_name_or_function: The device name or function to use in the context. |
| 5182 | |
| 5183 | Returns: |
| 5184 | A context manager that specifies the default device to use for newly |
| 5185 | created ops. |
| 5186 | |
| 5187 | Raises: |
| 5188 | RuntimeError: If eager execution is enabled and a function is passed in. |
| 5189 | """ |
| 5190 | if context.executing_eagerly(): |
| 5191 | if callable(device_name_or_function): |
| 5192 | raise RuntimeError( |
| 5193 | "tf.device does not support functions when eager execution " |
| 5194 | "is enabled.") |
| 5195 | return context.device(device_name_or_function) |
| 5196 | elif executing_eagerly_outside_functions(): |
| 5197 | @tf_contextlib.contextmanager |
| 5198 | def combined(device_name_or_function): |
| 5199 | with get_default_graph().device(device_name_or_function): |
| 5200 | if not callable(device_name_or_function): |
| 5201 | with context.device(device_name_or_function): |
| 5202 | yield |
| 5203 | else: |
| 5204 | yield |
| 5205 | return combined(device_name_or_function) |
| 5206 | else: |
| 5207 | return get_default_graph().device(device_name_or_function) |
| 5208 | |
| 5209 | |
| 5210 | @tf_export("device", v1=[]) |
no test coverage detected