def get_model_fn(model, train=False):
  """Create a function to give the output of the score-based model.

  Args:
    model: The score model.
    train: `True` for training and `False` for evaluation.

  Returns:
    A model function.
  """

  def model_fn(x, labels):
    """Compute the output of the score-based model.

    Args:
      x: A mini-batch of input data.
      labels: A mini-batch of conditioning variables for time steps. Should be interpreted differently
        for different models.

    Returns:
      A tuple of (model output, new mutable states)
    """
    # [Tag] 4. Model-Wrapper
    if not train:
      model.eval()
      # [Goto] 5. Model
      return model(x, labels)
    else:
      model.train()
      # [Goto] 5. Model
      return model(x, labels)

  return model_fn