Meghan 6ff4548b6e [AMP] Support XLA:TPU (#96370) ...

With https://github.com/pytorch/xla/pull/5148, https://github.com/pytorch/xla/pull/4740

With these changes
XLA:GPU users should use `torch.cuda.amp.autocast()` for AMP with float16
XLA:TPU users should use `torch.amp.autocast('xla')` for AMP with bfloat16

Pull Request resolved: https://github.com/pytorch/pytorch/pull/96370
Approved by: https://github.com/bdhirsh, https://github.com/malfet

2023-06-23 19:46:42 +00:00

..

__init__.py

add autocast cpu doc

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

autocast_mode.py

[AMP] Support XLA:TPU (#96370)

2023-06-23 19:46:42 +00:00