For future reference, iv’e managed to re-fold it, here is the solution I came up with:
#forward
start = torch.rand((2,3,64,64))
forward1 = start.unfold(2,16,16)
forward2 = forward1.unfold(3,16,16)
forward3 = forward2.flatten(start_dim=4, end_dim=5)
forward4 = forward3.flatten(start_dim=2, end_dim=3)
forward5 = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in forward4])


# going back
back1 = forward5.unfold(2,256,256).permute(0,2,1,3) 	# back1 == forward4
back2 = back1.unfold(2,4,4).permute(0,1,2,4,3) 			# back2 == forward3
back3 = back2.unfold(4,16,16) 							# back3 == forward2
back4 = back3.permute(0,1,2,3,5,4).reshape(2,3,4,-1,16) # back4 == forward1
back5 = back4.permute(0,1,2,4,3).reshape(2,3,-1,64)		# back5 == start

(back5 == start).all()

I don’t know if it is the most optimal way to do it, but it works