Minor fix to make sure printing does not crash if a filter_fn is used.

PiperOrigin-RevId: 371714419

Author: Mesh TensorFlow Team

mesh_tensorflow/transformer/utils.py

@@ -962,7 +962,8 @@ def serialized_fn(mtf_features):
           tf.logging.info("Variables in %s but not in graph:", init_checkpoint)
           tf.logging.info("\n".join(sorted(
               ckpt_vars -
-              {init_checkpoint_variable_mapping(v) for v in global_vars})))
+              {init_checkpoint_variable_mapping(v)
+               for v in filtered_global_vars})))
           tf.logging.info("Variables in graph but not in %s:", init_checkpoint)
           tf.logging.info("\n".join(sorted(global_vars - restore_vars)))
           tf.train.init_from_checkpoint(