Save the date! Google I/O returns May 18-20 Register now

tf_agents.utils.common.soft_device_placement

Context manager for soft device placement, allowing summaries on CPU.

Eager and graph contexts have different default device placements. See b/148408921 for details. This context manager should be used whenever using summary writers contexts to make sure summaries work when executing on TPUs.

Sets tf.config.set_soft_device_placement(True) within the context