maybe consider reshaping the tensors after unfolding, I did something like this once,
x = torch.randn(5, 5, 5, 5)
a = nn.Fold((5, 5), (3, 3))
d = x.unfold(2, 3, 1).unfold(3, 3, 1).reshape(x.size(0), -1, 9)
a(d).shape

torch.Size([5, 5, 5, 5])
to get the same shape, after unfolding then folding, pytorch version of fold, adds overlapping values.