Merge branch 'beam_search' into 'main'

fix a bug for size mismatch

See merge request ADLR/megatron-lm!438

megatron/text_generation/generation.py

@@ -395,7 +395,7 @@ def beam_search_and_return_on_first_stage(model, tokens, lengths, beam_size, sto
             # if cannot find stop token, add open beams to hyps
             if not done:
                 for beam_id in range(beam_size):
-                    beam_hyp.add(tokens[beam_id].clone(), scores[beam_id], context_length + 1 - prompt_length)
+                    beam_hyp.add(tokens[beam_id].clone(), scores[beam_id].squeeze(), context_length + 1 - prompt_length)
 
             # rank based on scores
             sorted_hyps = sorted(beam_hyp.beams, key=lambda x: x[0], reverse=True)