Added missing assignation

model.py

@@ -177,7 +177,7 @@ class VAE(tf.keras.Model):
     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)
+        x_input = 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))