Update app.py

app.py

@@ -38,8 +38,6 @@ def predict(req: Request):
     # 🔥 FIX HERE
     inputs.pop("token_type_ids", None)
 
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-
     with torch.no_grad():
         outputs = model(**inputs)
 
@@ -49,7 +47,7 @@ def predict(req: Request):
     pred = torch.argmax(probs, dim=1).item()
     confidence = probs[0][pred].item()
 
-    label = "ROADBLOCK" if pred == 1 else "NOT_ROADBLOCK"
+    label = model.config.id2label[pred]
 
     return {
         "input": req.text,