Fix NameError: name 'model_dict' is not defined

Closes #91
2018-12-31 13:29:39 +01:00 · 2018-12-31 13:29:39 +01:00 · 10524b885c
parent e78911bf88
commit 10524b885c
1 changed files with 16 additions and 13 deletions
--- a/train.py
+++ b/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 = {