def METHOD_NAME(_):
  # Build the train and eval datasets from the MNIST data. Also return the
  # input shape which is constructed based on the `image_data_format`
  # i.e channels_first or channels_last.
  tf.enable_eager_execution()
  train_ds, eval_ds, input_shape = get_input_datasets()
  # Instantiate the MirroredStrategy object. If we don't specify `num_gpus` or
  # the `devices` argument then all the GPUs available on the machine are used.
  # TODO(priyag): Use `tf.distribute.MirroredStrategy` once available.
  strategy = mirrored_strategy.MirroredStrategy(['/gpu:0', '/cpu:0'])
  # Create and compile the model under Distribution strategy scope.
  # `fit`, `evaluate` and `predict` will be distributed based on the strategy
  # model was compiled with.
  with strategy.scope():
    model = get_model(input_shape)
    optimizer = rmsprop.RMSProp(learning_rate=0.001)
    model.compile(loss=tf.keras.losses.categorical_crossentropy,
                  optimizer=optimizer,
                  metrics=['accuracy'])
  # Train the model with the train dataset.
  model.fit(x=train_ds, epochs=20, steps_per_epoch=468)
  # Evaluate the model with the eval dataset.
  score = model.evaluate(eval_ds, steps=10, verbose=0)
  print('Test loss:', score[0])
  print('Test accuracy:', score[1])