import tensorflow as tf

    """
    logits, losses = self(features)  # pylint: disable=not-callable
    if self.hparams.sampling_method == "argmax":
      samples = tf.argmax(logits, axis=-1)
    else:
      assert self.hparams.sampling_method == "random"
