I had to change this line to make it work correctly:
patches = patches.reshape(3, p_size, p_size, -1).permute(3, 0, 1, 2)
But now I get the following error message:
Traceback (most recent call last):
  File "fold.py", line 23, in <module>
    img_back = fold(patches)/fold(mask_p)
    ...
    raise NotImplementedError("Input Error: Only 3D input Tensors are supported (got {}D)".format(input.dim()))
NotImplementedError: Input Error: Only 3D input Tensors are supported (got 4D)

Specifically, running fold(patches) causes the error.