[pytorch] Expand PixelShuffle to support any number of batch dims (#49187)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49187
Expands the implementation of PixelShuffle to support any number of batch dimensions
Test Plan: `buck test caffe2/test:nn -- test_pixel_shuffle`
Reviewed By: mruberry
Differential Revision: D25399058
fbshipit-source-id: ab0a7f593b276cafc9ebb46a177e2c1dce56d0de
