diff --git a/optimum/tpu/generation/token_selector.py b/optimum/tpu/generation/token_selector.py
index e7cef097..97469f47 100644
--- a/optimum/tpu/generation/token_selector.py
+++ b/optimum/tpu/generation/token_selector.py
@@ -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.generation_config.get_generation_mode()
+        generation_mode = generation_config.get_generation_mode()
         if generation_mode not in [GenerationMode.GREEDY_SEARCH, GenerationMode.SAMPLE]:
             raise ValueError("Unsupported generation mode")