Use universal flatten APIs (#152505)
Summary:
X-link: https://github.com/pytorch/pytorch/pull/152505
Approved by: https://github.com/anijain2305
ghstack dependencies: #152389
Reviewed By: huydhn
Differential Revision: D74531500
fbshipit-source-id: ea17d89d5febe607ece2f43d053f5f533dbd1bd5
