Update app.py

app.py

@@ -33,15 +33,18 @@ def load_model():
 model, tokenizer = load_model()
 
 # Prediction function
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 def get_prediction(prompt):
     messages = [{"role": "user", "content": prompt}]
     inputs = tokenizer.apply_chat_template(
         messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
-    ).to("cuda")
+    ).to(device)
     output = model.generate(inputs, max_new_tokens=150, temperature=0.7, top_p=0.95)
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
     return decoded.split("###")[-1].strip()
 
+
 # UI Header
 st.title("🧠 AnthroBot")
 st.write("Enter your anthropometric estimates to receive an interpreted summary inputs — manually or via CSV upload.")