import tensorflow as tf

    x_means_hot, x_means, q_loss, e_loss = self.embedding_lookup(
        x_reshaped, self.means)

    if self.hparams.ema:
      tf.logging.info("Using EMA with beta = {}".format(self.hparams.beta))
      updated_ema_count = \
          moving_averages.assign_moving_average(
              self.ema_count,
              tf.reduce_sum(
                  tf.reshape(
                      x_means_hot,
                      shape=[-1, self.hparams.num_blocks,
                             self.hparams.block_v_size]),
                  axis=0),
              self.hparams.decay,
              zero_debias=False)

      dw = tf.matmul(
