diff --git a/aaa_helpers.txt b/aaa_helpers.txt
index 6273c2e..2086e8d 100644
--- a/aaa_helpers.txt
+++ b/aaa_helpers.txt
@@ -1,6 +1,7 @@
 nohup python auto_encoder_static.py --name="lhunor_ae_static_r7_3" > nohup/review/ae_static_r7_3.out 2>&1 &
 nohup python auto_encoder_body.py --name="lhunor_ae_body_r7_3" > nohup/review/ae_body_r7_3.out 2>&1 &
 nohup python auto_encoder_curv.py --name="lhunor_ae_curv_r7_3" > nohup/review/ae_curv_r7_3.out 2>&1 &
+nohup python cvae.py --name="lhunor_abl_all_resnet_r7" > nohup/review/abl_all_resnet_r7.out 2>&1 &
 
 one final convolutional layer for smoothing, on the final coordinates
 
diff --git a/configs/review/config_abl_all.py b/configs/review/config_abl_all.py
index 2ae6f4c..68985b5 100644
--- a/configs/review/config_abl_all.py
+++ b/configs/review/config_abl_all.py
@@ -19,8 +19,8 @@ config['only_curvatures'] =           False
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '0'
-config['name'] = 'abl_all_r7'
+config['gpu'] = '4'
+config['name'] = 'abl_all_resnet_r7'
 config['notes'] = ''
 
 config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
@@ -43,9 +43,9 @@ config['test_dir'] = '/home/workspace/data/dataset/test_ready'
 # config['frames_file_valid'] = 'frames_dress_30k.txt'
 # config['frames_file_test'] = 'frames_dress_30k.txt'
 
-config['static_model'] = 'checkpoints/post/static_r7_2'
-config['body_model'] = 'checkpoints/post/body_r7_2'
-config['curvature_model'] = 'checkpoints/post/curv_r7_2'
+config['static_model'] = 'checkpoints/post/static_r7_3'
+config['body_model'] = 'checkpoints/post/body_r7_3'
+config['curvature_model'] = 'checkpoints/post/curv_r7_3'
 
 
 config['smpl_res'] = 14475
diff --git a/cvae.py b/cvae.py
index cd46b2e..7691e02 100644
--- a/cvae.py
+++ b/cvae.py
@@ -58,7 +58,9 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
     for config_file_name in config_file_names:
-        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
@@ -73,7 +75,7 @@ def kl_loss(mean, logvar):
 
 # 3d loss on posed !!!!!!!!!!!!
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -474,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())
diff --git a/models/base.py b/models/base.py
index d0feffb..b070254 100644
--- a/models/base.py
+++ b/models/base.py
@@ -118,13 +118,29 @@ class VariationalEncoderNew(tf.keras.Model):
 
         self.encoder = resnet_18(features)
 
-        self.mean_branch = make_basic_block_layer(filter_num=features,
-                                             blocks=2,
-                                             stride=1)
+        self.mean_branch = tf.keras.Sequential([
+            tf.keras.layers.Conv2D(filters=features,
+                                            kernel_size=(3, 3),
+                                            strides=1,
+                                            padding="same"),
+            tf.keras.layers.BatchNormalization(),
+            tf.keras.layers.Conv2D(filters=features,
+                                            kernel_size=(3, 3),
+                                            strides=1,
+                                            padding="same")
+        ])
 
-        self.var_branch = make_basic_block_layer(filter_num=features,
-                                             blocks=2,
-                                             stride=1)
+        self.var_branch = tf.keras.Sequential([
+            tf.keras.layers.Conv2D(filters=features,
+                                            kernel_size=(3, 3),
+                                            strides=1,
+                                            padding="same"),
+            tf.keras.layers.BatchNormalization(),
+            tf.keras.layers.Conv2D(filters=features,
+                                            kernel_size=(3, 3),
+                                            strides=1,
+                                            padding="same")
+        ])
 
         self.flatten_mean = tf.keras.layers.Flatten()
         self.flatten_var = tf.keras.layers.Flatten()
diff --git a/models/cvae_model.py b/models/cvae_model.py
index ef2193f..195d93d 100644
--- a/models/cvae_model.py
+++ b/models/cvae_model.py
@@ -182,7 +182,7 @@ class CVAENew(tf.keras.Model):
         self.encoder = VariationalEncoderNew(config['cvae']['encoder_dims'][-1])
         self.decoder = DecoderNew(config['cvae']['decoder_dims'][-1], is_cvae=True)
 
-        self.reshape = tf.keras.layers.Reshape((16, 8, config['cvae']['decoder_dims'][-1]))
+        self.reshape = tf.keras.layers.Reshape((8, 4, config['cvae']['decoder_dims'][-1]))
 
         if config['avg_garment_path'] != '':
             self.avg_garment = np.load(config['avg_garment_path'])