FP16MomentumSGDUpdate Op fix and enable for ROCm (#15150)
Summary:
1. Fix a bug in FP16MomentumSGDUpdate operator
2. Enable operator for ROCm
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15150
Differential Revision: D13473145
Pulled By: bddppq
fbshipit-source-id: 4c5c5f30cb9bba658e3639dbe193fa08a304d306
