Just figured it out. Had to permute the patches:
patches = x.unfold(3, kernel_size, stride).unfold(2, kernel_size, stride)
becomes
patches = x.unfold(3, kernel_size, stride).unfold(2, kernel_size, stride).permute(0,1,2,3,5,4)