import torch


def make_onehot_kernel(kernel_size, index):
    """
    Make 2D one hot square kernel, i.e. h=w
    k[kernel_size, kernel_size] = 0 except k.view(-1)[index] = 1
    """
    kernel = torch.zeros(kernel_size, kernel_size)
    kernel.view(-1)[index] = 1
    return kernel.view(1, 1, kernel_size, kernel_size)