reword comments

The purpose of `input_size` is to be able to reshape the quantized output back to its original shape

Signed-off-by: reinvantveer <rein@vantveer.me>

examples/generative/vq_vae.py

@@ -93,9 +93,9 @@ def __init__(self, num_embeddings, embedding_dim, beta=0.25, **kwargs):
         )
 
     def call(self, x):
-        # Calculate the input shape of the inputs and
-        # then flatten the inputs keeping `embedding_dim` intact.
+        # Memorize the input shape of the inputs
         input_shape = tf.shape(x)
+        # Then flatten the inputs, keeping `embedding_dim` intact.
         flattened = tf.reshape(x, [-1, self.embedding_dim])
 
         # Quantization.