[optim] refactor adamax to use functional API (#55830)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/55830
ghstack-source-id: 126325537
Reviewed By: driazati
Differential Revision: D26561017
fbshipit-source-id: 41273d200e546d4ac08d39b57865d63c624f143a
