import tensorflow as tf

    # Reshape x_discrete
    shape_x = common_layers.shape_list(x)
    shape_discrete = shape_x[:-1]
    x_discrete = tf.reshape(x_discrete, shape_discrete)
    x_means = tf.reshape(x_means, shape=shape_x)
    h1 = x + tf.stop_gradient(x_means - x)

    h2 = tf.layers.dense(tf.nn.relu(h1), self.hparams.filter_size, name="vch2")
    res = tf.layers.dense(
        tf.nn.relu(h2), self.hparams.hidden_size, name="vcfin")
    embed_fn = partial(self.embed)
    return {
        "dense": res,
        "discrete": x_discrete,
