As you’ve explained, nn.Fold accumulated overlapping patches. If you want a custom operation, such as a masking with zeros, you could potentially create a mask from the patches, fold the tensors back, and mask the zeros spots afterwards with a multiplication.