Update app.py

app.py

@@ -19,7 +19,17 @@ class ToxicCommentTagger(pl.LightningModule):
         self.n_training_steps = n_training_steps
         self.n_warmup_steps = n_warmup_steps
         self.criterion = nn.BCELoss()
-
+        
+        
+    def forward(self, input_ids, attention_mask, labels=None):
+        output = self.bert(input_ids, attention_mask=attention_mask)
+        output = self.classifier(output.pooler_output)
+        output = torch.sigmoid(output)    
+        loss = 0
+        if labels is not None:
+            loss = self.criterion(output, labels)
+        return loss, output
+        
 
 def predict(model, tokenizer, sentence):