Return default precision that is supported by the hardware. Args: training: `-mixed` or `-true` version of the precision to use tpu: whether TPU device is used Returns: default precision that is suitable for the task and is supported by the hardware
(training: bool, tpu: bool = False)
| 489 | |
| 490 | |
| 491 | def get_default_supported_precision(training: bool, tpu: bool = False) -> str: |
| 492 | """Return default precision that is supported by the hardware. |
| 493 | |
| 494 | Args: |
| 495 | training: `-mixed` or `-true` version of the precision to use |
| 496 | tpu: whether TPU device is used |
| 497 | |
| 498 | Returns: |
| 499 | default precision that is suitable for the task and is supported by the hardware |
| 500 | """ |
| 501 | if tpu: |
| 502 | return "32-true" |
| 503 | if not torch.cuda.is_available() or torch.cuda.is_bf16_supported(): |
| 504 | return "bf16-mixed" if training else "bf16-true" |
| 505 | return "16-mixed" if training else "16-true" |