img_tensor: N, C, H, W

def add_depth_channel(img_tensor, pad_mode):
    '''
    img_tensor: N, C, H, W
    '''
    img_tensor[:, 1] = get_depth_tensor(pad_mode)
    img_tensor[:, 2] = img_tensor[:, 0] * get_depth_tensor(pad_mode)