Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Load model and tokenizer
 @st.cache_resource
@@ -34,9 +35,13 @@ def generate_prompt(comment):
 def get_response(comment):
     prompt = generate_prompt(comment)
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    
+    # Check if CUDA is available, otherwise use CPU
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     outputs = model.generate(
-        input_ids=inputs["input_ids"].to("cuda"), 
-        attention_mask=inputs["attention_mask"].to("cuda"), 
+        input_ids=inputs["input_ids"].to(device), 
+        attention_mask=inputs["attention_mask"].to(device), 
         max_new_tokens=140, 
         pad_token_id=tokenizer.pad_token_id  # Ensure padding is handled properly
     )