remove folder if error

2018-04-26 05:44:18 -07:00 · 2018-04-26 05:44:18 -07:00 · d143616ada
parent c894919120
commit d143616ada
1 changed files with 8 additions and 2 deletions
--- a/train.py
+++ b/train.py
@ -8,6 +8,7 @@ import signal
 import argparse
 import importlib
 import pickle
+import traceback
 import numpy as np

 import torch.nn as nn
@ -395,5 +396,10 @@ def main(args):


 if __name__ == '__main__':
-    signal.signal(signal.SIGINT, signal_handler)
-    main(args)
+    # signal.signal(signal.SIGINT, signal_handler)
+    try:
+        main(args)
+    except:
+        signal_handler()
+        traceback.print_exc()
+        sys.exit()