def complex_unpack(x):
    """
        Expects shape: b x 2c x height x width
        Returns shape: b x c x height x width x 2
    """
    b, c, h, w = x.shape
    assert c % 2 == 0
    newc = c//2
    return x.view(b, newc, 2, h, w).permute(0, 1, 3, 4, 2).contiguous()