import torch


def make_scores_mask(scores_shape, use_current_timestep=True, cuda=False):
    """
    Make upper triangular mask of 1s and 0s. If use_current_timestep is False,
    diagonal is also 1 (i.e. masked).
    """
    scores_mask = torch.ones(scores_shape, dtype=torch.uint8,
                             requires_grad=False)
    if cuda:
        scores_mask = scores_mask.cuda()

    scores_mask = torch.triu(scores_mask,
                             diagonal=1 if use_current_timestep else 0)
    return scores_mask