Update src/models/qa_model.py

src/models/qa_model.py

@@ -16,7 +16,7 @@ class QAEnsembleModel(nn.Module):
             model = AutoModelForQuestionAnswering.from_pretrained(model_name)#.half()
             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=0)
+                           tokenizer=model_name, device=device)
             self.nlps.append(nlp)
         self.entity_dict = entity_dict
         self.thr = thr