Update app.py

app.py

@@ -42,11 +42,12 @@ LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
 # "Qwen/Qwen2.5-7B-Instruct"
 
 llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+
 llm_model = AutoModelForCausalLM.from_pretrained(
     LLM_MODEL,
-    device_map="auto",
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
-)
+).to(device)
+
 llm_model.eval()
 
 # ============================================================