From 555715f41847d71d155dcfdaaefbdfa7084fff63 Mon Sep 17 00:00:00 2001
From: Yong Hoon Shin <48474650+sarckk@users.noreply.github.com>
Date: Tue, 20 May 2025 01:53:43 -0700
Subject: [PATCH] Fix broken example generation script for Llama3 (#38062)

Fix broken example generation script for llama3
---
 examples/pytorch/text-generation/run_generation.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/examples/pytorch/text-generation/run_generation.py b/examples/pytorch/text-generation/run_generation.py
index 9943a4f54a..4b50fbd07f 100755
--- a/examples/pytorch/text-generation/run_generation.py
+++ b/examples/pytorch/text-generation/run_generation.py
@@ -34,7 +34,6 @@ from transformers import (
     GPT2Tokenizer,
     GPTJForCausalLM,
     LlamaForCausalLM,
-    LlamaTokenizer,
     OpenAIGPTLMHeadModel,
     OpenAIGPTTokenizer,
     OPTForCausalLM,
@@ -63,7 +62,7 @@ MODEL_CLASSES = {
     "xlm": (XLMWithLMHeadModel, XLMTokenizer),
     "gptj": (GPTJForCausalLM, AutoTokenizer),
     "bloom": (BloomForCausalLM, BloomTokenizerFast),
-    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "llama": (LlamaForCausalLM, AutoTokenizer),
     "opt": (OPTForCausalLM, GPT2Tokenizer),
 }