Update app.py

app.py

@@ -17,10 +17,13 @@ def load_model(model_path):
         model_path,
         trust_remote_code=True,
         token=HF_TOKEN,
+        torch_dtype=torch.float16,  # 強制 FP16，避免 FP8 問題
         device_map="auto",  # 讓 transformers 自動決定使用 CPU/GPU
+        revision="main"
     )
     return model, tokenizer
 
+
 # 預設載入 DeepSeek-V3
 current_model, current_tokenizer = load_model("deepseek-ai/DeepSeek-V3")
 
@@ -33,7 +36,8 @@ def chat(message, history, model_name):
     if model_name != current_model:
         current_model, current_tokenizer = load_model(model_name)
 
-    inputs = current_tokenizer(message, return_tensors="pt").to("cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    inputs = current_tokenizer(message, return_tensors="pt").to(device)
     outputs = current_model.generate(**inputs, max_length=1024)
     response = current_tokenizer.decode(outputs[0], skip_special_tokens=True)