Update model.py

model.py

@@ -31,7 +31,7 @@ class CommentMTLModel(nn.Module):
         # Takes BERT's pooled output and maps it to toxicity logits (multi-label)
         self.toxicity_classifier = nn.Linear(self.bert.config.hidden_size, num_toxicity_labels)
 
-    def forward(self, input_ids, attention_mask):
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
         """
         Forward pass of the model.
 
@@ -48,6 +48,7 @@ class CommentMTLModel(nn.Module):
         outputs = self.bert(
             input_ids=input_ids,
             attention_mask=attention_mask
+            token_type_ids=token_type_ids
         )
 
         # Get the pooled output