small updates

scripts/ae_body_train.py

@@ -13,7 +13,8 @@ args = vars(parser.parse_args())
 
 config_module = importlib.import_module(args['config'])
 config = config_module.config
-# from configs.ae_body_cape_m import config
+if args['debug']:
+    from configs.ae_body_cape_m import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:

scripts/cvae.py

@@ -12,10 +12,12 @@ args = vars(parser.parse_args())
 
 config_module = importlib.import_module(args['config'])
 config = config_module.config
-
+if args['debug']:
+    pass
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
-# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import time
 import argparse
@@ -488,7 +490,7 @@ if __name__ == '__main__':
         step = load_model(model, config['pretrained_model'])
 
     if not args['debug']:
-        model = torch.compile(model)
+        # model = torch.compile(model)
         wandb.watch(model, log='all', log_freq=1000, log_graph=True)
-        
+
     train()

scripts/upscaler_train.py

@@ -38,7 +38,7 @@ from utils.utils import InfiniteDataLoader as DataLoader
 from pytorch_msssim import MS_SSIM
 
 
-low_res = 32
+low_res = 64
 is_skirt = False
 is_cape = True
 is_male = True