code: def train_mnist(config):
      # https://github.com/tensorflow/tensorflow/issues/32159
      import tensorflow as tf
  
      batch_size = 128
      num_classes = 10
      epochs = 12
  
      with FileLock(os.path.expanduser("~/.data.lock")):
          (x_train, y_train), (x_test, y_test) = mnist.load_data()
      x_train, x_test = x_train / 255.0, x_test / 255.0
      model = tf.keras.models.Sequential(
          [
              tf.keras.layers.Flatten(input_shape=(28, 28)),
              tf.keras.layers.Dense(
                  config["hidden"], activation="relu"
              ),
              tf.keras.layers.Dropout(0.2),
              tf.keras.layers.Dense(num_classes, activation="softmax"),
          ]
      )
  
      model.compile(
          loss="sparse_categorical_crossentropy",
          optimizer=tf.keras.optimizers.SGD(
              lr=config["lr"], momentum=config["momentum"]
          ),
          metrics=["accuracy"],
      )
  
      model.fit(
          x_train,
          y_train,
          batch_size=batch_size,
          epochs=epochs,
          verbose=0,
          validation_data=(x_test, y_test),
          callbacks=[
              ReportCheckpointCallback(
                  checkpoint_on=[],
                  metrics={"mean_accuracy": "accuracy"},
              )
          ],
      )
  
file_path: ray/python/ray/tune/examples/tune_mnist_keras.py
imports: import argparse
  import os
  from filelock import FileLock
  from tensorflow.keras.datasets import mnist
  import ray
  from ray import train, tune
  from ray.tune.schedulers import AsyncHyperBandScheduler
  from ray.air.integrations.keras import ReportCheckpointCallback
label_desc: Use image node_id 1 for calling a global function with example usage: train_mnist(config) without return types