Update app.py

app.py

@@ -7,11 +7,12 @@ import gradio as gr
 peft_model_id = "Pr123/TinyLlama-EA-Chat"
 
 # Load Model with PEFT adapter
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoPeftModelForCausalLM.from_pretrained(
     peft_model_id,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
+    torch_dtype=torch.float16  # Keeping half precision
+).to(device)
+
 tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=500)