NVIDIA / apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch

BSD 3-Clause "New" or "Revised" License

8.42k stars 1.4k forks source link

Closed eqy closed 10 months ago

eqy commented 11 months ago

1654

I think the tests pass without the changes to multi_tensor_apply.cuh, but leaving it as-is makes me a bit nervous...

TODOs: graph-capturable Adam, and all other optimizers if people really need 64-bit indexing there...

CC @crcrpar