Update app.py

app.py

@@ -67,9 +67,9 @@ class ToxicModel(torch.nn.Module):
 def inference_fn(model, input_ids=None, attention_mask=None, token_type_ids=None):  
     model.eval()
     print(input_ids, attention_mask, token_type_ids)
-    input_ids = input_ids.to(device)  
-    attention_mask = attention_mask.to(device)  
-    token_type_ids = token_type_ids.to(device)  
+    input_ids = input_ids[0].to(device)  
+    attention_mask = attention_mask[0].to(device)  
+    token_type_ids = token_type_ids[0].to(device)  
     
     with torch.no_grad():
         output = model(input_ids.unsqueeze(0), attention_mask.unsqueeze(0), token_type_ids.unsqueeze(0))