Update app.py

app.py

@@ -65,6 +65,7 @@ class ToxicModel(torch.nn.Module):
  
 def inference_fn(model, input_ids=None, attention_mask=None, token_type_ids=None):  
     model.eval()
+    
     input_ids = input_ids.to(device)  
     attention_mask = attention_mask.to(device)  
     token_type_ids = token_type_ids.to(device)  
@@ -91,6 +92,10 @@ def predict(comment=None) -> dict:
     mask = inputs['attention_mask']
     token_type_ids = inputs["token_type_ids"]
     
+    ids = torch.tensor(ids, dtype=torch.long),
+    mask = torch.tensor(mask, dtype=torch.long),
+    token_type_ids = torch.tensor(token_type_ids, dtype=torch.long),
+    
     model = ToxicModel()
     
     model.load_state_dict(torch.load("toxicx_model_0.pth", map_location=torch.device(device)))