H = KERNEL_SIZE[0]
W = KERNEL_SIZE[1]
N = OUT_CHANNEL
M = IN_CHANNEL

# implement context mask
mask = np.zeros((N, M, H, W))
mask[:,:,:(H//2)] = 1
mask[:,:,(H//2),range((W//2))] = 1

# implement channel mask
if mask_type == 'A':
    mask[np.triu_indices(N, k=0, m=M)] = 0
elif mask_type == 'B':
    mask[np.triu_indices(N, k=1, m=M)] = 0