def gen_3D_flat(yz1_grid, depth1, depth2):
    """
    depth1/2 are B*1*N, yz1_grid is 3*N(N=H*W)
    transform both to B*C(1 or 3)*N
    both outputs are B*3*N
    """
    batch_size = depth1.shape[0]

    depth1_flat = depth1.expand(-1, 3, -1) # B*3*N
    depth2_flat = depth2.expand(-1, 3, -1)
    yz1_grid_batch = yz1_grid.expand(batch_size, -1, -1) # B*3*N
    
    xyz_1 = yz1_grid_batch * depth1_flat
    xyz_2 = yz1_grid_batch * depth2_flat

    return xyz_1, xyz_2