fix NEL overfitting test for GPU (#5236)

This commit is contained in:
Sofie Van Landeghem 2020-04-02 10:32:52 +02:00 committed by GitHub
parent 311133e579
commit ab59f3124e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1456,7 +1456,7 @@ class EntityLinker(Pipe):
scores = prior_probs + sims - (prior_probs*sims)
# TODO: thresholding
best_index = scores.argmax()
best_index = scores.argmax().item()
best_candidate = candidates[best_index]
final_kb_ids.append(best_candidate.entity_)
final_tensors.append(sentence_encoding)