From 10989715d001f2f735614f5db019452c71cc18d5 Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 9 Mar 2020 20:25:09 +0100 Subject: [PATCH] rename variable --- src/transformers/modeling_tf_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py index 2ab1a8a1c03..cfaa10b0467 100644 --- a/src/transformers/modeling_tf_utils.py +++ b/src/transformers/modeling_tf_utils.py @@ -990,7 +990,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin): next_scores, (batch_size, num_beams * vocab_size) ) # (batch_size, num_beams * vocab_size) - next_scores, next_tokens = tf.math.top_k(next_scores, 2 * num_beams, sorted=True) + next_scores, next_tokens = tf.math.top_k(next_scores, k=2 * num_beams, sorted=True) assert shape_list(next_scores) == shape_list(next_tokens) == [batch_size, 2 * num_beams]