Update app.py

app.py

@@ -98,7 +98,7 @@ DEFAULT_CONTENT = "Put your content sample here."
 # Function to make predictions
 def predict(content, policy):
     input_text = PROMPT.format(policy=policy, content=content)
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
 
     with torch.inference_mode():
         outputs = model(input_ids)