Update src/models/qa_model.py

src/models/qa_model.py

@@ -9,12 +9,12 @@ from src.features.graph_utils import find_best_cluster
 class QAEnsembleModel(nn.Module):
 
     def __init__(self, model_name, model_checkpoints, entity_dict,
-                 thr=0.1, device="cuda:0"):
+                 thr=0.1, device="cpu"):
         super(QAEnsembleModel, self).__init__()
         self.nlps = []
         for model_checkpoint in model_checkpoints:
             model = AutoModelForQuestionAnswering.from_pretrained(model_name).half()
-            model.load_state_dict(torch.load(model_checkpoint), strict=False)
+            model.load_state_dict(torch.load(model_checkpoint, map_location=torch.device('cpu')), strict=False)
             nlp = pipeline('question-answering', model=model,
                            tokenizer=model_name, device=int(device.split(":")[-1]))
             self.nlps.append(nlp)