Handle 2D blocking with foreach (#107840)
Previously blocking in foreach ops was only 1D. This PR allows handling kernels with 2D blocking with foreach as well.
Code when at least one dim matches:
[example code + output](https://gist.github.com/mlazos/9f100b21cfe2540f0a24303a8349c196)
Code when neither X or Y dim matches:
[example code + output](https://gist.github.com/mlazos/14e2a455f635896dface09be601595dd)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/107840
Approved by: https://github.com/jansel
