import torch


def subsequent_mask(lens):
    """Mask out future word
    Args:
        lens: (bs,)
    Return:
        mask: (bs, max_len, max_len)
    """
    bs, max_len = len(lens), max(lens)
    mask = torch.ones([bs, max_len, max_len]).tril_(0)
    mask = mask > 0
    return mask