tf_agents.utils.common.soft_device_placement

View source on GitHub

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.

Yields:

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