From b240317d761f1a85c4acbaec38003846576d95d9 Mon Sep 17 00:00:00 2001 From: Anthony Wang Date: Mon, 21 Feb 2022 16:39:58 -0600 Subject: [PATCH] Take top 5 instead of top 3 --- predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/predict.py b/predict.py index 53447b9..9ef453d 100644 --- a/predict.py +++ b/predict.py @@ -2,7 +2,7 @@ import numpy as np import torch -def predict(device, dataset, model, text, next_words=100, top_k=3): +def predict(device, dataset, model, text, next_words=100, top_k=5): model.eval() words = text.split()