diff --git a/deeptrack/models/vaes/vae.py b/deeptrack/models/vaes/vae.py
index e1f10dbff..25c6e72ae 100644
--- a/deeptrack/models/vaes/vae.py
+++ b/deeptrack/models/vaes/vae.py
@@ -30,7 +30,7 @@ def train_step(self, data):
 
             # Sample a random point in the latent space
             epsilon = tf.random.normal(shape=tf.shape(z_mean))
-            z = z_mean + tf.exp(z_log_var) * epsilon
+            z = z_mean + tf.exp(0.5 * z_log_var) * epsilon
 
             # Reconstruct the input image
             rdata = self.decoder(z)