Update app.py

app.py

@@ -4,9 +4,12 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the model and tokenizer
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2-medium")
+model.to(device)
+
 tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2-medium", clean_up_tokenization_spaces=True)
 
-model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2-medium")
 
 # Define the initial prompt for the system
 system_prompt = """