code: def get_compiled_model(lr=0.001):
      from tensorflow import keras
  
      model = keras.models.Sequential(
          [
              keras.layers.Flatten(input_shape=(28, 28)),
              keras.layers.Dense(128, activation="relu"),
              keras.layers.Dropout(0.2),
              keras.layers.Dense(10),
          ]
      )
  
      model.compile(
          optimizer=keras.optimizers.Adam(learning_rate=lr),
          loss=keras.losses.SparseCategoricalCrossentropy(
              from_logits=True
          ),
          metrics=["accuracy"],
      )
  
      return model
  
imports: import logging
  from pyspark.sql import SparkSession
  from examples.spark_dataset_converter.utils import get_mnist_dir
  from petastorm.spark import SparkDatasetConverter, make_spark_converter
label_desc: Use image node_id 1 for calling a global function with example usage: get_compiled_model(lr) and returns: model