class Sampling(layers.Layer):

    """Uses (mean, log_var) to sample z, the vector encoding a digit."""

    def call(self, inputs):

        mean, log_var = inputs

        batch = tf.shape(mean)[0]

        dim = tf.shape(mean)[1]

        return mean + tf.exp(0.5 * log_var) * epsilon