I have function that translates image into puzzles using unfold:
def make_jigsaw_puzzle(x, grid_size=(2, 2)):
    # x shape is C x H x W
    C, H, W = x.size()

    assert H % grid_size[0] == 0
    assert W % grid_size[1] == 0

    C, H, W = x.size()
    x_jigsaw = x.unfold(1, H // grid_size[0], W // grid_size[1])
    x_jigsaw = x_jigsaw.unfold(2, H // grid_size[0], W // grid_size[1])
    x_jigsaw = x_jigsaw.contiguous().view(-1, C,  H // grid_size[0], W // grid_size[1])
    return x_jigsaw

x_jigsaw shape is grid_size[0] * grid_size[1] x C x H x W.
How can I get image from patches back?