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?