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