Fix qa predict device setting.
@@ -563,7 +563,7 @@ class AnswerExtractor(Transformer):
if cache_model:
self.save_model()
- def predict(self, test_dataloader, num_gpus=None, local_rank=-1, verbose=True):
+ def predict(self, test_dataloader, num_gpus=None, verbose=True):
"""
Predicts answer start and end logits.
@@ -584,11 +584,16 @@ class AnswerExtractor(Transformer):
def _to_list(tensor):
return tensor.detach().cpu().tolist()
- device, num_gpus = get_device(num_gpus=num_gpus, local_rank=local_rank)
+ device, num_gpus = get_device(num_gpus=num_gpus, local_rank=-1)
- self.model.to(device)
+ if num_gpus > 1:
+ if not isinstance(self.model, torch.nn.DataParallel):
+ self.model = torch.nn.DataParallel(self.model)
+ else:
+ if isinstance(self.model, torch.nn.DataParallel):
+ self.model = self.model.module
- # score
+ self.model.to(device)
self.model.eval()
all_results = []
