def pad_sentences(sents, pad_id):
    """
    Adding pad_id to sentences in a mini-batch to ensure that 
    all augmented sentences in a mini-batch have the same word length.
    Args:
        sents: list(list(int)), a list of a list of word ids
        pad_id: the word id of the "<pad>" token
    Return:
        aug_sents: list(list(int)), |s_1| == |s_i|, for s_i in sents
    """

    max_seq_length = 0

    for sentence in sents:
        if len(sentence)>max_seq_length:
            max_seq_length = len(sentence)
    
    for i in range(len(sents)):
        if len(sents[i])<max_seq_length:
            sents[i]+=[pad_id]*(max_seq_length-len(sents[i]))
            #print(sents[i])
    return sents