From d143616ada05fe69651cc3e898367a8389d422d4 Mon Sep 17 00:00:00 2001
From: Eren Golge <egolge@mozilla.com>
Date: Thu, 26 Apr 2018 05:44:18 -0700
Subject: [PATCH] remove folder if error

---
 train.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index 7e41712a..54ec06ef 100644
--- 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()