small updates

2023-05-22 00:14:18 +02:00
parent 52e461ca47
commit 6b81c0d1a6
5 changed files with 202 additions and 7 deletions
@@ -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()