diff --git a/configs/abl_all.py b/configs/abl_all.py
index 972b78a..e4458b5 100644
--- a/configs/abl_all.py
+++ b/configs/abl_all.py
@@ -2,12 +2,13 @@ from utils.convnext_utils import CNBlockConfig
 
 config = dict()
 
-config['gpu'] = '0'
-config['name'] = 'abl_torch_all'
+config['gpu'] = '6'
+config['name'] = 'abl_torch_all_2'
 config['notes'] = ''
 config['tags'] = ['iccv', 'abl_all', 'torch']
 config['script'] = 'scripts/cvae.py'
 
+config['pretrained_model'] = 'checkpoints/iccv/abl_torch_all'
 
 config['use_skinning'] =              True # and deshaping
 config['remove_global_orientation'] = True
@@ -43,15 +44,12 @@ config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
 config['body_model'] = 'checkpoints/iccv/body_torch_1'
 config['curvature_model'] = 'checkpoints/iccv/curv_torch_1'
 
-
 config['smpl_res'] = 14475
 if config['is_cape']:
   config['smpl_res'] = 3638
 
 config['latent_dim'] = 128
 
-config['pretrained_model'] = ''
-
 ###############    
 config['lr'] = 0.0001#                 0.0001
 ###############
diff --git a/configs/ae_curv_1.py b/configs/ae_curv_1.py
index fa8ff62..df9a25a 100644
--- a/configs/ae_curv_1.py
+++ b/configs/ae_curv_1.py
@@ -53,7 +53,6 @@ block_setting_reverse = [
     CNBlockConfig(64, 32, 9),
     CNBlockConfig(32, 16, 3),
     CNBlockConfig(16, 3, 3),
-    # CNBlockConfig(3, 3, 3),
 ]
 
 config['cond_ae']['encoder_block_settings'] = block_setting
diff --git a/configs/ae_static_orig.py b/configs/ae_static_orig.py
index 5a6080a..8180389 100644
--- a/configs/ae_static_orig.py
+++ b/configs/ae_static_orig.py
@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_1'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_1'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'
@@ -60,7 +62,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',
diff --git a/configs/ae_static_orig_2.py b/configs/ae_static_orig_2.py
index c02c91a..f47cba1 100644
--- a/configs/ae_static_orig_2.py
+++ b/configs/ae_static_orig_2.py
@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_2'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_2'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'
diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py
index 6986bcc..ee693f8 100644
--- a/scripts/ae_static_train.py
+++ b/scripts/ae_static_train.py
@@ -2,7 +2,16 @@ import os
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from configs.ae_static import config
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
@@ -98,7 +107,9 @@ def valid_step(batch):
        
 
 def plot_everything(batch: dict, type: str):
-    pred = model([batch['uv_static'], batch['uv_mask']]).detach()
+    model.eval()
+    with torch.inference_mode():
+        pred = model([batch['uv_static'], batch['uv_mask']]).detach()
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
@@ -162,9 +173,9 @@ if __name__ == '__main__':
     init_wandb()
     # device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
-    args = vars(parser.parse_args())
+    # parser = argparse.ArgumentParser(description='')
+    # parser.add_argument('--name', type=str)
+    # args = vars(parser.parse_args())
     
     resolution_modifier = ''
     if config['resolution'] != 512: