import torch


def neighbours_mask(size, window_size):
    """Mask for neighbour positions.
    Args:
        size (int): squared tensor size
        window_size (int): how many elements to be considered as valid around
            the ith element (including ith).
    Returns:
        torch.Tensor: (size, size)
    """
    z = torch.ones(size, size, dtype=torch.uint8)
    mask = (torch.triu(z, diagonal=1 + window_size // 2)
            + torch.tril(z, diagonal=- window_size // 2))
    return z - mask