code: def train_and_save(models, x_train, y_train, x_test, y_test):
      for model in models:
          model.compile(
              loss=tf.keras.losses.categorical_crossentropy,
              optimizer="adam",
              metrics=["accuracy"],
          )
  
          # Print the model summary.
          model.summary()
  
          # Add a pruning step callback to peg the pruning step to the optimizer's
          # step. Also add a callback to add pruning summaries to tensorboard
          callbacks = [
              pruning_callbacks.UpdatePruningStep(),
              pruning_callbacks.PruningSummaries(
                  log_dir=FLAGS.output_dir
              ),
          ]
  
          model.fit(
              x_train,
              y_train,
              batch_size=batch_size,
              epochs=epochs,
              verbose=1,
              callbacks=callbacks,
              validation_data=(x_test, y_test),
          )
          score = model.evaluate(x_test, y_test, verbose=0)
          print("Test loss:", score[0])
          print("Test accuracy:", score[1])
  
          # Export and import the model. Check that accuracy persists.
          saved_model_dir = "/tmp/saved_model"
          print("Saving model to: ", saved_model_dir)
          tf.keras.models.save_model(
              model, saved_model_dir, save_format="tf"
          )
          print("Loading model from: ", saved_model_dir)
          loaded_model = tf.keras.models.load_model(saved_model_dir)
  
          score = loaded_model.evaluate(x_test, y_test, verbose=0)
          print("Test loss:", score[0])
          print("Test accuracy:", score[1])
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_cnn.py
imports: from __future__ import print_function
  from absl import app
  from absl import flags
  import tensorflow
  from tensorflow_model_optimization.python.core.sparsity.keras import prune
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
label_desc: Use image node_id 4 for calling a global function with example usage: train_and_save(models, x_train, y_train, x_test, y_test) without return types