def forward(self, batch):
    return self.model(input_ids=batch['diff_input_ids'], attention_mask=batch['diff_attention_mask'], decoder_input_ids=batch['msg_input_ids'], decoder_attention_mask=batch['msg_attention_mask'], labels=batch['msg_labels'])
def forward(self, batch):
    return self.model(input_ids=batch['diff_input_ids'], attention_mask=batch['diff_attention_mask'], decoder_input_ids=batch['msg_input_ids'], decoder_attention_mask=batch['msg_attention_mask'], labels=batch['msg_labels'])