Fix scalar summaries with TF2 import

PiperOrigin-RevId: 324257955

Author: William Fedus

mesh_tensorflow/utils.py

@@ -124,7 +124,7 @@ def host_call_fn(global_step, *args):
       assert len(args) == len(summaries)
       for i, tensor in enumerate(args):
         name = summaries[i][0]
-        tf.summary.scalar(
+        tf2.summary.scalar(
             name, tf.reduce_mean(tensor), step=global_step)
       return tf.summary.all_v2_summary_ops()