Optimize alignBatchDimsAtFront (#41941)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/41941
If we know that the tensor already has the desired aligned size, we
don't need to put in the effort to align it.
Test Plan: - `./build/bin/vmap_test`, `pytest test/test_vmap.py -v`
Reviewed By: albanD
Differential Revision: D22764101
Pulled By: zou3519
fbshipit-source-id: a2ab7ce7b98d405ae905f7fd98db097210bfad65
