Update train.py

tools/train.py

@@ -6,8 +6,8 @@ import torch
 from mmcv import Config
 
 from mmdet import __version__
-from mmcv.runner import init_dist
-from mmdet.apis import (get_root_logger, set_random_seed, train_detector)
+from mmdet.apis import (get_root_logger, init_dist, set_random_seed,
+                        train_detector)
 from mmdet.datasets import build_dataset
 from mmdet.models import build_detector
 
@@ -97,20 +97,14 @@ def main():
             CLASSES=datasets[0].CLASSES)
     # add an attribute for visualization convenience
     model.CLASSES = datasets[0].CLASSES
-    
-    try:
-        train_detector(
-            model,
-            datasets,
-            cfg,
-            distributed=distributed,
-            validate=args.validate,
-            logger=logger)
-
-    except Exception as error:
-        print("Error... saving checkpoint")
-        torch.save(model.state_dict(), './models/crash.pth')
-        raise error
-        
+    train_detector(
+        model,
+        datasets,
+        cfg,
+        distributed=distributed,
+        validate=args.validate,
+        logger=logger)
+
+
 if __name__ == '__main__':
     main()