No public description

PiperOrigin-RevId: 592656146

Author: chaoyan1037

official/projects/yt8m/modeling/yt8m_model_utils.py

@@ -120,7 +120,9 @@ def frame_pooling(frames, method="average", num_frames=None):
           tf.ones_like(frames, dtype=frames.dtype)
           * _large_compatible_negative(frames.dtype),
       )
-    reduced = tf.reduce_max(frames, 1)
+    # Magic to avoid loss NaN when bfloat16 is enabled.
+    # See yaqs/5377152819545505792 and b/214396297 for more discussion.
+    reduced = tf.reduce_max(frames, 1) + tf.reduce_mean(frames, 1) * 0
   elif method == "swap":
     # Note we assume the frames are in the shape of
     # [batch_size, num_frames, feature_size]. Otherwise this function might