import tensorflow as tf

      updated_ema_means = updated_ema_means / tf.expand_dims(
          updated_ema_count, axis=-1)

      with tf.control_dependencies([e_loss]):
        update_means = tf.assign(self.means, updated_ema_means)
        with tf.control_dependencies([update_means]):
          loss += self.hparams.beta * e_loss
    else:
      # Use a gradient based loss for learning the cluster centers
      loss += q_loss + self.hparams.beta * e_loss
