import tensorflow as tf

    return logits

  def top(self, body_output, features):
    if isinstance(body_output, dict):
      logits = {}
      for k, v in body_output.iteritems():
        logits[k] = self._top_single(v, features)
      return logits
    else:
      return self._top_single(body_output, features)

  def _loss_single(self, logits, features):
    if not self._problem_hparams:
      log_warn(_no_problem_err("loss"))
      return (tf.constant(0., dtype=tf.float32),
              tf.constant(1., dtype=tf.float32))

    target_modality = self._problem_hparams.target_modality
    loss_num, loss_den = target_modality.loss(logits, features["targets"])
    loss_num *= self._problem_hparams.loss_multiplier
    return loss_num, loss_den

  def loss(self, logits, features):
    if isinstance(logits, dict):
      losses = {}
      for k, v in logits.iteritems():
        losses[k] = self._loss_single(v, features)
      return tf.add_n([n / d for n, d in logits.values()])
    else:
      return self._loss_single(logits, features)
