Update app.py

app.py

@@ -20,18 +20,16 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True, trust_remote_
 try:
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        device_map="auto",
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        trust_remote_code=True,
+        torch_dtype="auto",
+        device_map="cuda",
     )
 except Exception as e:
     print("Automatic device_map load failed, falling back to cpu. Error:", e)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        device_map={"": "cpu"},
-        torch_dtype=torch.float32,
-        trust_remote_code=True,
-    )
+        torch_dtype="auto",
+        device_map="auto",
+        )
 
 model.eval()
 print("Model loaded. Device:", next(model.parameters()).device)