fix model

app.py

@@ -23,11 +23,12 @@ transform = transforms.Compose([transforms.ToTensor(),
 
 # model imports
 def load_pretrained_params(model, model_state_path: str):
-    pretrained_dict = torch.load(model_state_path, map_location="cpu")
+    checkpoint = torch.load(model_state_path, map_location="cpu")
+    pretrained_dict = checkpoint["state_dict"]
     model_dict = model.state_dict()
     # 1. filter out unnecessary keys
-    if list(pretrained_dict.keys())[0].startswith("module."):
-        pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
+    if list(pretrained_dict.keys())[0].startswith("model."):
+        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}
     else:
         pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
     # 2. overwrite entries in the existing state dict