Roy Hvaara bc69a19139 [MPS] Add support for bf16 autocast (#139390) ...

This PR adds support for bf16 autocast. Most of the code and ideas are copied from #99272.

Most of the heavy lifting was done by AI.

Fixes #139386

Pull Request resolved: https://github.com/pytorch/pytorch/pull/139390
Approved by: https://github.com/malfet

Co-authored-by: Kulin Seth <kulin_seth@apple.com>
Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>

2024-11-20 19:52:28 +00:00

..

__init__.py

generalize custom_fwd&custom_bwd to be device-agnostic (#126531)

2024-05-25 06:48:16 +00:00

autocast_mode.py

[MPS] Add support for bf16 autocast (#139390)

2024-11-20 19:52:28 +00:00

grad_scaler.py

Add support to GradScaler for respecting an already set grad_scale value (#123429)

2024-06-27 22:40:54 +00:00