def METHOD_NAME(features, labels, mode, params):
  """Function specifying the model that is required by the `tf.estimator` API.
  Args:
    features: Input images
    labels: Labels of images
    mode: One of `ModeKeys.TRAIN`, `ModeKeys.EVAL` or 'ModeKeys.PREDICT'
    params: A dictionary of extra parameter that might be passed
  Returns:
    An instance of `tf.estimator.EstimatorSpec`
  """
  inputs = features
  if isinstance(inputs, dict):
    inputs = features["image"]
  config = params["config"]
  model = revnet.RevNet(config=config)
  if mode == tf.estimator.ModeKeys.TRAIN:
    global_step = tf.train.get_or_create_global_step()
    learning_rate = tf.train.piecewise_constant(
        global_step, config.lr_decay_steps, config.lr_list)
    optimizer = tf.train.MomentumOptimizer(
        learning_rate, momentum=config.momentum)
    logits, saved_hidden = model(inputs, training=True)
    grads, loss = model.compute_gradients(saved_hidden, labels, training=True)
    with tf.control_dependencies(model.get_updates_for(inputs)):
      train_op = optimizer.apply_gradients(
          zip(grads, model.trainable_variables), global_step=global_step)
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
  else:
    logits, _ = model(inputs, training=False)
    predictions = tf.argmax(logits, axis=1)
    probabilities = tf.nn.softmax(logits)
    if mode == tf.estimator.ModeKeys.EVAL:
      loss = model.compute_loss(labels=labels, logits=logits)
      return tf.estimator.EstimatorSpec(
          mode=mode,
          loss=loss,
          eval_metric_ops={
              "accuracy":
                  tf.metrics.accuracy(labels=labels, predictions=predictions)
          })
    else:  # mode == tf.estimator.ModeKeys.PREDICT
      result = {
          "classes": predictions,
          "probabilities": probabilities,
      }
      return tf.estimator.EstimatorSpec(
          mode=mode,
          predictions=predictions,
          export_outputs={
              "classify": tf.estimator.export.PredictOutput(result)
          })