[FSDP] Improved shape unflattening test (#72573)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72573
Verify shapes are restored appropriately in this test.
ghstack-source-id: 149111043
Test Plan: CI
Reviewed By: zhaojuanmao
Differential Revision: D34101125
fbshipit-source-id: 94260da2b7420cf58c5569e596885aa65fe7726e
(cherry picked from commit e57a30e8e4caea0593836e52084194a3d3497b72)
