Update app.py

app.py

@@ -5,12 +5,13 @@ from peft import PeftModel
 import torch
 
 # Load model and tokenizer once using caching
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 @st.cache_resource
 def load_model():
     base_model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-1B-Instruct")
     model = PeftModel.from_pretrained(base_model, "CallmeKaito/llama-3.2-1b-it-brainrot")
     tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-1B-Instruct")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
     model = model.to(device)
     return model, tokenizer
 
@@ -51,7 +52,7 @@ if prompt := st.chat_input("What's up?"):
         )
         
         # Tokenize and generate
-        inputs = tokenizer(chat_prompt, return_tensors="pt").to("cuda")
+        inputs = tokenizer(chat_prompt, return_tensors="pt").to(device)
         outputs = model.generate(
             **inputs,
             max_new_tokens=500,