Context manager for mixed precision training. Args: enabled: Whether to enable autocast. dtype: torch.float16 or torch.bfloat16. If None, defaults to float16.
(enabled, dtype=None)
| 23 | |
| 24 | @contextmanager |
| 25 | def maybe_autocast(enabled, dtype=None): |
| 26 | """Context manager for mixed precision training. |
| 27 | |
| 28 | Args: |
| 29 | enabled: Whether to enable autocast. |
| 30 | dtype: torch.float16 or torch.bfloat16. If None, defaults to float16. |
| 31 | """ |
| 32 | if enabled: |
| 33 | with autocast(dtype=dtype): |
| 34 | yield |
| 35 | else: |
| 36 | yield |
| 37 | |
| 38 | |
| 39 | class Trainer: |
no test coverage detected
searching dependent graphs…