diff --git a/aaa_helpers.txt b/aaa_helpers.txt
index f8ea6ae..964044f 100644
--- a/aaa_helpers.txt
+++ b/aaa_helpers.txt
@@ -1,3 +1,5 @@
+nohup python cvae.py --name="lhunor_abl_all_10p_curv_10_r2" > nohup/review/abl_all_10p_curv_10_r2.out 2>&1 &
+
 one final convolutional layer for smoothing, on the final coordinates
 
 skirts!!!!
diff --git a/configs/config_abl_all_10p_curv_10_r2.py b/configs/config_abl_all_10p_curv_10_r2.py
index 710053a..762eff8 100644
--- a/configs/config_abl_all_10p_curv_10_r2.py
+++ b/configs/config_abl_all_10p_curv_10_r2.py
@@ -57,6 +57,7 @@ if config['is_cape']:
 
 
 config['pretrained_model'] = ''
+config['latent_dim'] = 128
 
 ###############    
 config['lr'] = 0.0001#                 0.0001
@@ -88,42 +89,37 @@ config['use_one_hot'] = False
 
 config['cvae'] = dict()
 config['cvae']['is_label'] = False
+config['cvae']['latent_dim'] = config['latent_dim']
 config['cvae']['add_noise'] = True
-config['cvae']['noise_amount'] = 0.001
+config['cvae']['noise_amount'] = 0.025
 
 config['cvae']['kernel_size'] = 7
-config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3]
-config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64]
-config['cvae']['decoder_depths'] = config['cvae']['encoder_depths']
-config['cvae']['decoder_dims'] = config['cvae']['encoder_dims']
+config['cvae']['encoder_depths'] = [3, 3, 9, 3]
+config['cvae']['encoder_dims'] = [32, 64, 128, 256]
+config['cvae']['decoder_input'] = (16, 8, 80)
+config['cvae']['decoder_depths'] = [3, 3, 9, 3]
+config['cvae']['decoder_dims'] = [256, 128, 64, 16]
 
 
-config['cond_ae_static'] = dict()
-config['cond_ae_static']['is_label'] = False
-config['cond_ae_static']['kernel_size'] = 7
-config['cond_ae_static']['bottleneck_size'] = 8
-config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
-config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
-config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
-config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
-config['cond_ae_static']['activation'] = 'tanh'
-
-config['cond_ae_curv'] = dict()
-config['cond_ae_curv']['is_label'] = True
-config['cond_ae_curv']['kernel_size'] = 7
-config['cond_ae_curv']['bottleneck_size'] = 40
-config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
-config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
-config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
-config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
-config['cond_ae_curv']['activation'] = 'tanh'
+config['cond_ae'] = dict()
+config['cond_ae']['is_label'] = False
+config['cond_ae']['upscale_output'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 64
+config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3]
+config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottleneck_size']]
+config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
+config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
+config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
 
 config['cond_ae_body'] = dict()
-config['cond_ae_body']['is_label'] = False
+config['cond_ae_body']['latent_dim'] = config['latent_dim']
 config['cond_ae_body']['kernel_size'] = 7
-config['cond_ae_body']['bottleneck_size'] = 16
-config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
-config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
-config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
-config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
-config['cond_ae_body']['activation'] = 'tanh'
\ No newline at end of file
+config['cond_ae_body']['upscale_output'] = True
+config['cond_ae_body']['bottleneck_size'] = 512
+config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]
diff --git a/cvae.py b/cvae.py
index d25acf6..21c8ffa 100644
--- a/cvae.py
+++ b/cvae.py
@@ -41,7 +41,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAENew as CVAE
+from models.cvae_model import CVAE
 
 
 def init_wandb():
@@ -73,7 +73,7 @@ def kl_loss(mean, logvar):
 
 # 3d loss on posed !!!!!!!!!!!!
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -476,7 +476,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())