Update app.py

app.py

@@ -18,11 +18,12 @@ tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=False)
 tokenizer.pad_token = tokenizer.eos_token
 
 def predict_risk(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = {k: v.to("cpu") for k, v in inputs.items()}
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=50,      # keeps output fast but complete
+            max_new_tokens=50,
             temperature=0.7,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id