Update main.py

main.py

@@ -9,6 +9,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 #Load pre-trained tokenizer and model (Works)
 model_name = "gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="auto",
@@ -52,7 +53,7 @@ async def generate_text(item: Item):
 
     # logging.info("Response generated")
 
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True).to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True, ).to(model.device)
 
 
     # input_ids = tokenizer.encode(item.prompt, return_tensors="pt")