transformers-community
/

group-beam-search

Text Generation

custom_generate

text-generation-inference

Model card Files Files and versions

manueldeprada HF Staff commited on Aug 27

Commit

eff00e5

·

verified ·

1 Parent(s): b60e744

Upload folder using huggingface_hub

Files changed (1) hide show

custom_generate/generate.py +1 -1

custom_generate/generate.py CHANGED Viewed

@@ -208,7 +208,7 @@ def _group_beam_search(
     # define beam scorer
     beam_scorer = BeamSearchScorer(
-        batch_size=input_ids.shape[0],
         num_beams=generation_config.num_beams,
         device=input_ids.device,
         length_penalty=generation_config.length_penalty,

     # define beam scorer
     beam_scorer = BeamSearchScorer(
+        batch_size=input_ids.shape[0] // generation_config.num_beams,
         num_beams=generation_config.num_beams,
         device=input_ids.device,
         length_penalty=generation_config.length_penalty,