-
Ruoxin Sang authored
When calling `strategy.reduce` in eager mode, wrap the `strategy.run` calls inside with `tf.function` so it is compatible with TPUStrategy. PiperOrigin-RevId: 312847673 Change-Id: I6db92c34ba24e160689da3fca2fe0a3c26223d52
4797b3b9