diff --git a/models/cvae_norm_model.py b/models/cvae_norm_model.py
index 327f5fb..e1b8596 100644
--- a/models/cvae_norm_model.py
+++ b/models/cvae_norm_model.py
@@ -36,27 +36,16 @@ class CVAENorm(nn.Module):
         self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
 
 
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+    def sample(self, uv_template, uv_body_posed, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
+            eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
         else:
             eps = noise
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
-        if add_noise:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-        
-        z = tf.concat([eps, z_cond], axis=-1)
-
-        x = self.decoder(z, training=False)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
-        if self.use_offsets_over_body:
-            x = x + uv_body_posed
-
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        z_cond = self.encode_conds(uv_template, uv_body_posed)
+        z = torch.cat([eps, z_cond], dim=1)
+        x = self.decoder(z)
+        x = x * uv_mask[:,None,:,:]
 
         return x