Skip to content

Commit

Permalink
fix: get_generation_mode is now a method of generation_config
Browse files Browse the repository at this point in the history
API change when transformers was updated.
  • Loading branch information
tengomucho committed Apr 10, 2024
1 parent fe888a9 commit dbf11f7
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion optimum/tpu/generation/token_selector.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ def create(
logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
generation_config.pad_token_id = eos_token_id

generation_mode = model._get_generation_mode(generation_config, None)
generation_mode = model.generation_config.get_generation_mode()
if generation_mode not in [GenerationMode.GREEDY_SEARCH, GenerationMode.SAMPLE]:
raise ValueError("Unsupported generation mode")

Expand Down

0 comments on commit dbf11f7

Please sign in to comment.