Update app.py

app.py

@@ -35,13 +35,14 @@ model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
 text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
 
 # Load the Llama-3 model and tokenizer once during startup
+device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer_llama = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", token=hf_token)
 model_llama = AutoModelForCausalLM.from_pretrained(
     "meta-llama/Meta-Llama-3-8B",
     torch_dtype='float16',
     device_map='auto',
     token=hf_token
-)
+).to(device) 
 
 # Define your prompt template
 prompt_template = """..."""  # Your existing prompt template here