fix ByteTensor -> LongTensor to avoid overflow
@@ -155,7 +155,7 @@ class EntityRankingJob(EvaluationJob):
def _get_rank(self, scores, answers):
answers = answers.reshape((-1, 1)).expand(-1, self.dataset.num_entities).long()
true_scores = torch.gather(scores, 1, answers)
- ranks = torch.sum(scores > true_scores, dim=1)
+ ranks = torch.sum((scores > true_scores).long(), dim=1)
return ranks
def _get_metrics(self, rank_hist, suffix=''):
