File tree Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change @@ -222,7 +222,7 @@ def process_labels(labels, x):
222222beam_size = 1
223223
224224model_input = transform (input_text )
225- model_output = sequence_generator .generate (model_input , eos_idx = eos_idx , beam_size = beam_size )
225+ model_output = sequence_generator .generate (model_input , eos_idx = eos_idx , num_beams = beam_size )
226226output_text = transform .decode (model_output .tolist ())
227227
228228for i in range (cnndm_batch_size ):
@@ -312,7 +312,7 @@ def process_labels(labels, x):
312312beam_size = 1
313313
314314model_input = transform (input_text )
315- model_output = sequence_generator .generate (model_input , eos_idx = eos_idx , beam_size = beam_size )
315+ model_output = sequence_generator .generate (model_input , eos_idx = eos_idx , num_beams = beam_size )
316316output_text = transform .decode (model_output .tolist ())
317317
318318for i in range (imdb_batch_size ):
You can’t perform that action at this time.
0 commit comments