64bit indexing fused adam (#5187)
## The Issue
Applying `FusedAdam` on large tensors will cause an error `CUDA error:
an illegal memory access was encountered`.
https://github.com/microsoft/DeepSpeed/issues/3429
https://github.com/NVIDIA/apex/issues/1654
## PR Content
Following the solution in the apex repository
(https://github.com/NVIDIA/apex/pull/1765), changing indexing type to
`int64` if necessary.
---------
Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
