import torch


def subsequent_mask(size, device):
    """Mask out subsequent positions (adapted from
    http://nlp.seas.harvard.edu/2018/04/03/attention.html)"""
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape, device=device), diagonal=1)
    return subsequent_mask == 0