Fixed dimensionality issues in ``encode``

model.py

@@ -174,7 +174,11 @@ class VAE(tf.keras.Model):
     sd: tf.Tensor
         The standard deviation parameter of the distribution.
     """
-      
+    x_input = tf.expand_dims(x_input, axis=-1) # Add channel dimension
+
+    if tf.rank(x_input) == 3: # If there's no batch dimension, add it
+        tf.expand_dims(x_input, axis=0)
+
     mu, rho = tf.split(self.encoder(x_input), num_or_size_splits=2, axis=1)
     sd = tf.math.log(1 + tf.math.exp(rho))
     z_sample = mu + sd * tf.random.normal(shape=(120,))