Get autocast context manager based on precision.
(precision: str, device_type: str = 'cuda')
| 326 | |
| 327 | |
| 328 | def get_autocast_context(precision: str, device_type: str = 'cuda'): |
| 329 | """Get autocast context manager based on precision.""" |
| 330 | from contextlib import nullcontext |
| 331 | if precision in ('bf16', 'bfloat16'): |
| 332 | return partial(torch.amp.autocast, device_type=device_type, dtype=torch.bfloat16) |
| 333 | elif precision in ('fp16', 'float16'): |
| 334 | return partial(torch.amp.autocast, device_type=device_type, dtype=torch.float16) |
| 335 | return nullcontext |
| 336 | |
| 337 | |
| 338 | # ============================================================================ |