Update app.py

app.py

@@ -34,7 +34,7 @@ def load_model(model_name: str):
     metadata_features = 0
     N_UNIQUE_CLASSES = 38  
     
-    if model_name == 'gena-bert':
+    if model_name == 'GENA-Bert':
         base_model = AutoModel.from_pretrained('AIRI-Institute/gena-lm-bert-base-lastln-t2t', trust_remote_code=True, output_hidden_states=True)
         tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bert-base-lastln-t2t', trust_remote_code=True)
     
@@ -80,7 +80,7 @@ def load_model(model_name: str):
             model_seq,
             new_head
         )
-        weights = torch.load('CNN_1stGEAC_m2_best.pth')
+        weights = torch.load('CNN_1stGEAC_m2_best.pth',map_location=torch.device('cpu'))
         model.load_state_dict(weights)
         return model, None
 
@@ -112,7 +112,7 @@ def load_model(model_name: str):
             new_head
         )
         
-        weights = torch.load('NOO_CNN_1stGEAC_m4_16kcw_best.pth')
+        weights = torch.load('NOO_CNN_1stGEAC_m4_16kcw_best.pth',map_location=torch.device('cpu'))
         joined_model.load_state_dict(weights)
         
         return joined_model, None