save

models/pyramid_level.py

@@ -39,22 +39,16 @@ class PyramidLevel(nn.Module):
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = torch.randn(z_cond.shape).to(uv_mask.device)
         else:
             eps = noise
 
-        z = torch.concat([eps, z_cond], dim=-1)
-
+        z = torch.cat([eps, z_cond], dim=1)
         x = self.decoder(z)
-        
-        if self.avg_garment is not None:
-            x = x + self.avg_garment[None,:,:,:]
-        
-        x = x * uv_mask[:,None:,:]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
-
     def reparameterize(self, mean, logvar, fix_noise=False):
         if self.training == True:
             if not fix_noise: