Update gpt_example.py

Fix bug of tf gpt example
NVIDIA · Mar 26, 2022 · ba6960e · ba6960e
1 parent f8f0ee8
commit ba6960e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/tensorflow/gpt/gpt_example.py b/examples/tensorflow/gpt/gpt_example.py
@@ -236,6 +236,7 @@ def ft_gpt_op(var_dict,
 
     output_ids, parent_ids, sequence_length, cum_log_probs = gpt_op_module.gpt(
         extended_input_ids, # 0
+        extended_input_lengths, # 1
         [tf.cast(var_dict["model/h%d/ln_1/b:0" % l], data_type) for l in range(decoder_args.num_layer)], # 2
         [tf.cast(var_dict["model/h%d/ln_1/g:0" % l], data_type) for l in range(decoder_args.num_layer)], # 3
         [tf.cast(var_dict["model/h%d/attn/c_attn/w:0" % l], data_type) for l in range(decoder_args.num_layer)], # 4