Merge pull request #92 from yweweler/yweweler-fix-model-dict-restore

Fix NameError: name 'model_dict' is not defined

train.py

@@ -400,7 +400,10 @@ def main(args):
 
     if args.restore_path:
         checkpoint = torch.load(args.restore_path)
+        try:
             model.load_state_dict(checkpoint['model'])
+        except:
+            model_dict = model.state_dict()
             # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
             # 1. filter out unnecessary keys
             pretrained_dict = {