import torch


def block_eye(n, size):
    """
    Create a block_diagonal matrix of n blocks, where each block
    is torch.eye(size)
    """
    m1 = torch.ones(n, size, 1, size)
    m2 = torch.eye(n).view(n, 1, n, 1)
    return (m1*m2).view(n*size, n*size).to(torch.uint8)