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: jeanschmidt
Differential Revision: D74668908
fbshipit-source-id: 80df9f6c05096bf042f364a9412f8c60f6cf291f
