code: def setup_model(input_shape, image_train, label_train):
      """Baseline model."""
      model = tf.keras.Sequential(
          [
              tf.keras.layers.InputLayer(input_shape),
              tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
              tf.keras.layers.Conv2D(
                  filters=12, kernel_size=(3, 3), activation=tf.nn.relu
              ),
              tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
              tf.keras.layers.Flatten(),
              tf.keras.layers.Dense(10),
          ]
      )
      compile_and_fit(model, image_train, label_train, 5)
  
      return model
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/cluster_preserve_qat/keras/mnist_cnn.py
imports: from __future__ import print_function
  import os
  from absl import app
  import numpy
  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.quantization.keras import quantize
  from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.cluster_preserve import default_8bit_cluster_preserve_quantize_scheme
  from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.cluster_preserve.cluster_utils import strip_clustering_cqat
label_desc: Use image node_id 1 for calling a global function with example usage: setup_model(input_shape, image_train, label_train) and returns: model