diff --git a/configs/config_abl_all_post.py b/configs/config_abl_all_post.py
index 5093733..a00421b 100644
--- a/configs/config_abl_all_post.py
+++ b/configs/config_abl_all_post.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'] = '1'
-config['name'] = 'al_all_post'
+config['gpu'] = '0'
+config['name'] = 'abl_all_post_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
diff --git a/cvae.py b/cvae.py
index 59d7c75..e9aa853 100644
--- a/cvae.py
+++ b/cvae.py
@@ -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())
diff --git a/evaluate_ae.py b/evaluate_ae.py
index 28bb889..fdd4575 100644
--- a/evaluate_ae.py
+++ b/evaluate_ae.py
@@ -102,7 +102,7 @@ if __name__ == '__main__':
     ids = [id.split('_')[0] for id in ids]
     ids = set(ids)
 
-    iteration = ''
+    iteration = '70000'
 
     load_model(model, 'checkpoints/post/' + checkpoint_name, iteration=iteration)
     # load_model(model, 'checkpoints/' + 'abl_all_120k')