code: def _build_model():
      i = tf.keras.layers.Input(shape=(28, 28), name="input")
      x = tf.keras.layers.Reshape((28, 28, 1))(i)
      x = tf.keras.layers.Conv2D(
          20, 5, activation="relu", padding="valid", name="conv1"
      )(x)
      x = tf.keras.layers.MaxPool2D(2, 2)(x)
      x = tf.keras.layers.Conv2D(
          50, 5, activation="relu", padding="valid", name="conv2"
      )(x)
      x = tf.keras.layers.MaxPool2D(2, 2)(x)
      x = tf.keras.layers.Flatten()(x)
      x = tf.keras.layers.Dense(500, activation="relu", name="fc1")(x)
      output = tf.keras.layers.Dense(10, name="fc2")(x)
  
      model = tf.keras.Model(inputs=[i], outputs=[output])
      return model
  
file_path: model-optimization/tensorflow_model_optimization/python/core/common/keras/compression/algorithms/different_training_and_inference_test.py
imports: import os
  import tempfile
  import tensorflow
  from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import different_training_and_inference
  from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
label_desc: Use image node_id 1 for calling a global function with example usage: _build_model() and returns: model