Update app.py

app.py

@@ -66,8 +66,6 @@ class CNN_linear(nn.Module):
         self.fc = nn.Linear(in_features = embed_dim, out_features = self.nodes)
         self.linear = nn.Linear(in_features = self.nbr_filters, out_features = self.nodes)
         self.output = nn.Linear(in_features = self.nodes, out_features = 1)
-        self.direct_output = nn.Linear(in_features = embed_dim, out_features = 1)
-        self.magic_output = nn.Linear(in_features = 1, out_features = 1)
             
     def forward(self, tokens, need_head_weights=True, return_contacts=False, return_representation=True):
         
@@ -151,8 +149,8 @@ def predict_raw(raw_input):
 
     model = CNN_linear()
     st.write(model.state_dict().keys())
-    st.write(torch.load(modelfile, map_location=torch.device('cpu')).keys())
-    model.esm2.load_state_dict({k.replace('module.', ''):v for k,v in torch.load(modelfile, map_location=torch.device('cpu')).items()}, strict = False)
+    st.write({k.replace('module.', ''):v for k,v in torch.load(modelfile, map_location=torch.device('cpu')).items()}.keys())
+    model.load_state_dict({k.replace('module.', ''):v for k,v in torch.load(modelfile, map_location=torch.device('cpu')).items()}, strict = True)
     # model.load_state_dict(torch.load(modelfile, map_location=torch.device('cpu')), strict = False)
 
     # st.write('====Predict====')