修改模型load

Models.py

@@ -93,15 +93,15 @@ def get_model_token_classification(opt, src_vocab, trg_vocab):
        
     if opt.load_weights is not None:
         print("loading pretrained weights...")
-        model.load_state_dict(torch.load(f'{opt.load_weights}'))
-        # checkpoint = torch.load(opt.load_weights, map_location=lambda storage, loc: storage)
-        # state_dict = {}
-        # for k, v in checkpoint.items():
-        #     if k.startswith('module'):
-        #         state_dict[k[7:]] = v
-        #     else:
-        #         state_dict[k] = v
-        # model.load_state_dict(state_dict)
+        # model.load_state_dict(torch.load(f'{opt.load_weights}'))
+        checkpoint = torch.load(opt.load_weights, map_location=lambda storage, loc: storage)
+        state_dict = {}
+        for k, v in checkpoint.items():
+            if k.startswith('module'):
+                state_dict[k[7:]] = v
+            else:
+                state_dict[k] = v
+        model.load_state_dict(state_dict)
     else:
         for p in model.parameters():
             if p.dim() > 1: