[Checkpoint] Fix autocasting (#81766)
Add support for the correct autocasting in the non-reentrant checkpoint as it exists in the reentrant-version.
This was noticed by @awgu.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81766
Approved by: https://github.com/albanD
