code: def get_model(**kwargs):
      model = keras.Sequential(
          [
              keras.Input(shape=input_shape),
              layers.Conv2D(
                  kwargs["conv1"], kernel_size=(3, 3), activation="relu"
              ),
              layers.MaxPooling2D(pool_size=(2, 2)),
              layers.Conv2D(
                  kwargs["conv2"], kernel_size=(3, 3), activation="relu"
              ),
              layers.MaxPooling2D(pool_size=(2, 2)),
              layers.Flatten(),
              layers.Dropout(kwargs["dropout"]),
              layers.Dense(num_classes, activation="softmax"),
          ]
      )
  
      return model
  
file_path: auptimizer/Examples/compression/mnist_tensorflow/mnist.py
imports: from aup import print_result, aup_args
  import tensorflow
  import numpy
  from tensorflow import keras
  from tensorflow.keras import layers
  import tqdm
  from math import log
  import sys
  import aup
label_desc: Use image node_id 1 for calling a global function with example usage: get_model() and returns: model