ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

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