run training by changing just run.py

scripts/ae_body_train.py

@@ -2,13 +2,22 @@ import os
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from configs.ae_body_1 import config
+import argparse
+import importlib
+
+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
+# from configs.ae_body_1 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
 
 import time
-import argparse
 import wandb
 import random
 import numpy as np
@@ -159,12 +168,12 @@ def train():
 
 
 if __name__ == '__main__':
-    init_wandb()
+    # 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: