refactor autocast to be extensible for devices (#57104)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/57104
Test Plan: Imported from OSS
Reviewed By: gchanan
Differential Revision: D28094173
Pulled By: ezyang
fbshipit-source-id: a5fb62b9a4e58f30d2756bba4331d5fc88136b89
