configs

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: