Update app.py

app.py

@@ -1,12 +1,15 @@
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("TuringsSolutions/Gemma2LegalEdition", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("TuringsSolutions/Gemma2LegalEdition", trust_remote_code=True)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 def predict(prompt, temperature, max_tokens):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = model.generate(
         **inputs,
         max_new_tokens=max_tokens,