tensor.unfold should yield the same result:
a = torch.randn(64, 3, 16, 16)
kernel_size = 4
kernel_stride = 4
a = a.unfold(2, kernel_size, kernel_stride).unfold(3, kernel_size, kernel_stride)
a = a.contiguous().view(a.size(0), a.size(1), -1, a.size(4), a.size(5))
print(a.shape)
> torch.Size([64, 3, 16, 4, 4])

Not that the patches are now in dim2. You can of course permute the tensor as you wish.
chunks[0] now corresponds to a[:, :, 0].