Update src/models/predict_model.py

src/models/predict_model.py

@@ -23,11 +23,11 @@ title2idx = dict([(x.strip(), y) for x, y in zip(df_wiki.title, df_wiki.index.va
 
 qa_model = QAEnsembleModel("nguyenvulebinh/vi-mrc-large", ["models/qa_model_robust.bin"], entity_dict)
 pairwise_model_stage1 = PairwiseModel("nguyenvulebinh/vi-mrc-base")#.half()
-pairwise_model_stage1.load_state_dict(torch.load("models/pairwise_v2.bin"))
+pairwise_model_stage1.load_state_dict(torch.load("models/pairwise_v2.bin", map_location=torch.device('cpu')))
 pairwise_model_stage1.eval()
 
 pairwise_model_stage2 = PairwiseModel("nguyenvulebinh/vi-mrc-base")#.half()
-pairwise_model_stage2.load_state_dict(torch.load("models/pairwise_stage2_seed0.bin"))
+pairwise_model_stage2.load_state_dict(torch.load("models/pairwise_stage2_seed0.bin", map_location=torch.device('cpu')))
 
 bm25_model_stage1 = BM25Gensim("models/bm25_stage1/", entity_dict, title2idx)
 bm25_model_stage2_full = BM25Gensim("models/bm25_stage2/full_text/", entity_dict, title2idx)