shibo d03799f9a5 optimize the AMP func name in custom_device_mod (#98052) ...

Fixes #ISSUE_NUMBER
1、optimize the func name of AMP in custom device module，use `torch.foo.set_autocast_enable` instead of `torch.foo.set_autocast_foo_enable`.
2、In AMP with custom device，use `custom_device_mod.set_autocast_enable` instead of `getattr(custom_device_mod,  "set_autocast_enable"`, because we have check that `custom_device_mod` hasattr `set_autocast_enable` before.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/98052
Approved by: https://github.com/bdhirsh

2023-03-31 17:04:32 +00:00

..

__init__.py

add autocast cpu doc

2022-03-22 02:02:43 +00:00

autocast_mode.py

optimize the AMP func name in custom_device_mod (#98052)

2023-03-31 17:04:32 +00:00