code: def train_model(model, x_train, y_train, x_test, y_test):
      model.compile(
          loss=tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=True
          ),
          optimizer="adam",
          metrics=["accuracy"],
      )
  
      # Print the model summary.
      model.summary()
  
      # Model needs to be clustered after initial training
      # and having achieved good accuracy
      model.fit(
          x_train,
          y_train,
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          validation_split=0.1,
      )
  
      score = model.evaluate(x_test, y_test, verbose=0)
      print("Test loss:", score[0])
      print("Test accuracy:", score[1])
  
      return model
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/clustering/keras/mnist/mnist_cnn.py
imports: from __future__ import print_function
  import datetime
  import os
  from absl import app
  from absl import flags
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import cluster
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
  from tensorflow_model_optimization.python.core.clustering.keras import clustering_callbacks
label_desc: Use image node_id 3 for calling a global function with example usage: train_model(model, x_train, y_train, x_test, y_test) and returns: model