def METHOD_NAME(self, batch):
    input_ids_padded = []
    for input_ids, input_lens in batch: # input_ids are returned by this dataset (see __getitem__)
        pad_len = self.pad_max - input_ids.shape[0]
        #input_ids = F.pad(input_ids, pad=(0, pad_size), value=self.tokenizer.pad_token_id)
        input_ids_padded.append(input_ids)
    input_ids_padded = torch.vstack(input_ids_padded)
    return (input_ids_padded, input_ids_padded)