No public description

PiperOrigin-RevId: 771154434

Author: sagunb

orbit/utils/loop_fns.py

@@ -19,6 +19,10 @@
 
 import tensorflow as tf, tf_keras
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.tpu.google.sparse_core import tpu_embedding_v3
+# pylint: enable=g-direct-tensorflow-import
+
 
 def create_loop_fn(step_fn):
   """Creates a loop function driven by a Python `while` loop.
@@ -200,7 +204,8 @@ class LoopFnWithSummaries(tpu_summaries.OptionalSummariesFunction):
 
   def __call__(self, iterator, num_steps):
     if tf.summary.should_record_summaries():
-      output = self.with_summaries(iterator, tf.constant(1))
+      with tpu_embedding_v3.SequentialEmbeddingContext():
+        output = self.with_summaries(iterator, tf.constant(1))
       num_steps -= 1
     if num_steps >= 1:
       output = self.without_summaries(iterator, num_steps)