def METHOD_NAME(features, labels, mode, params):
  model = params['model']
  feature_columns = params['feature_columns']
  head = params['head']
  if mode == model_fn_lib.ModeKeys.TRAIN:
    logits = model.build_model(features, feature_columns, is_training=True)
  elif mode == model_fn_lib.ModeKeys.EVAL:
    logits = model.build_model(features, feature_columns, is_training=False)
  else:
    raise NotImplementedError
  def _train_op_fn(loss):
    global_step = training_util.get_global_step()
    assert global_step
    train_step = model.get_train_step(loss)
    with ops.control_dependencies(train_step):
      with ops.get_default_graph().colocate_with(global_step):
        return state_ops.assign_add(global_step, 1).op
  return head.create_model_fn_ops(
      features=features,
      mode=mode,
      labels=labels,
      train_op_fn=_train_op_fn,
      logits=logits)