diff options
Diffstat (limited to 'predict.py')
-rw-r--r-- | predict.py | 2 |
1 files changed, 1 insertions, 1 deletions
@@ -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() |