import torch


def create_mask(seq_len, device):
    """
    Helper function to create a square mask
    
    Parameters
    ----------
    seq_len : int
        length of the edge

    Returns
    -------
    2d tensor of the following shape
    0 1 1 ... 1
    0 0 1 ... 1
    ...
    0 0 0 ... 0
    """
    ones = torch.ones(seq_len, seq_len)
    mask = torch.triu(ones, diagonal=1).bool().to(device)
    
    return mask