import tensorflow as tf

        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

    # Get the discrete latent representation
    x_means_idx = tf.argmax(x_means_hot, axis=-1)

    # Get the binary representation
    num_bits = int(self.hparams.z_size // self.hparams.num_blocks)
    x_means_bits = self.int_to_bit(x_means_idx, num_bits=num_bits, base=2)
    x_discrete = self.bit_to_int(
        tf.to_int32(x_means_bits), num_bits=self.hparams.z_size, base=2)
