Context-manager to force placement of operations and Tensors on a device. Example: ```python with tf.device('gpu:0'): with tf.device('cpu:0'): shape = tf.constant([], dtype=tf.int32) x = tf.random.truncated_normal(shape, tf.float32) ``` will ensure that the `shape` Tensor is
(name)
| 1657 | |
| 1658 | |
| 1659 | def device(name): |
| 1660 | """Context-manager to force placement of operations and Tensors on a device. |
| 1661 | |
| 1662 | Example: |
| 1663 | ```python |
| 1664 | with tf.device('gpu:0'): |
| 1665 | with tf.device('cpu:0'): |
| 1666 | shape = tf.constant([], dtype=tf.int32) |
| 1667 | x = tf.random.truncated_normal(shape, tf.float32) |
| 1668 | ``` |
| 1669 | will ensure that the `shape` Tensor is on CPU but the `truncated_normal` |
| 1670 | operation runs on GPU 0. |
| 1671 | |
| 1672 | Args: |
| 1673 | name: Name of the device (see context().devices()), or None to |
| 1674 | perform automatic placement. |
| 1675 | |
| 1676 | Returns: |
| 1677 | Context manager for setting the device. |
| 1678 | """ |
| 1679 | ensure_initialized() |
| 1680 | return context().device(name) |
| 1681 | |
| 1682 | |
| 1683 | @tf_export("config.experimental_list_devices") |
nothing calls this directly
no test coverage detected