save

models/pyramid_level.py

@@ -101,6 +101,7 @@ class PyramidLevelNew(tf.keras.Model):
     def __init__(self, config):
         super(PyramidLevelNew, self).__init__()
         self.size = config['resolution']
+        self.latent_dim=config['latent_dim']
 
         self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim'])
         self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
@@ -126,13 +127,12 @@ class PyramidLevelNew(tf.keras.Model):
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = tf.random.normal(shape=(uv_mask.shape[0], 8, 4, self.latent_dim))
         else:
             eps = noise
 
         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,:,:,:]