Update src/models/qa_model.py

src/models/qa_model.py

@@ -9,7 +9,7 @@ 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="cpu"):
+                 thr=0.1, device="CPU:0"):
         super(QAEnsembleModel, self).__init__()
         self.nlps = []
         for model_checkpoint in model_checkpoints: