Fix: Load CPU-safe model for HF Space

app.py

@@ -13,11 +13,9 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 # Force CPU-safe model loading
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
-    device_map="cpu",                # Ensure model loads on CPU
-    torch_dtype=torch.float32        # Ensure safe dtype for CPU
-)
-
-device = torch.device("cpu")         # Explicitly define CPU
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True
+).to("cpu")
 
 
 # Dictionaries to decode user inputs
@@ -44,7 +42,7 @@ def get_prediction(age, gender, height, weight, ap_hi, ap_lo,
 
 Diagnosis:"""
 
-    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    inputs = tokenizer(input_text, return_tensors="pt").to("cpu")
     model.eval()
     with torch.no_grad():
         outputs = model.generate(**inputs, max_new_tokens=4)