diff --git a/models/cvae_model.py b/models/cvae_model.py
index 0009d5b..cfff553 100644
--- a/models/cvae_model.py
+++ b/models/cvae_model.py
@@ -49,22 +49,22 @@ class CVAE(tf.keras.Model):
             self.avg_garment = None
 
 
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
+    @tf.function
+    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None):
         if noise is None:
             eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
         else:
             eps = noise
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
         z = tf.concat([eps, z_cond], axis=-1)
 
-        x = self.decoder(z)
+        x = self.decoder(z, training=False)
         
         if self.avg_garment is not None:
             x = x + self.avg_garment[tf.newaxis,:,:,:]
         
-        if uv_mask is not None:
-            x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,:,:,tf.newaxis]
 
         return x
 
@@ -78,13 +78,13 @@ class CVAE(tf.keras.Model):
         return eps * tf.exp(logvar * .5) + mean
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
-        z_body = self.body_enc(uv_body_posed)
-        z_static = self.static_enc(uv_template)
-        z_curv = self.curv_enc(uv_curv)
+    def encode_conds(self, uv_template, uv_body_posed, uv_curv, training=None):
+        z_body = self.body_enc(uv_body_posed, training=training)
+        z_static = self.static_enc(uv_template, training=training)
+        z_curv = self.curv_enc(uv_curv, training=training)
 
         z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
-        z_cond = self.cond_encoder(z_cond)
+        z_cond = self.cond_encoder(z_cond, training=training)
 
         return z_cond
 
@@ -101,7 +101,7 @@ class CVAE(tf.keras.Model):
         x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
         mean, logvar = self.encoder(x, training=training)
         z = self.reparameterize(mean, logvar)
-        z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high)
+        z_cond = self.encode_conds(uv_template_high, uv_body_posed_high, uv_curv_high, training=training)
 
         z = tf.concat([z, z_cond], axis=-1)
         x = self.decoder(z, training=training)