diff --git a/pytorch_lightning/plugins/precision/native_amp.py b/pytorch_lightning/plugins/precision/native_amp.py index ae9f261085229..4c34bed972067 100644 --- a/pytorch_lightning/plugins/precision/native_amp.py +++ b/pytorch_lightning/plugins/precision/native_amp.py @@ -19,7 +19,7 @@ import pytorch_lightning as pl from pytorch_lightning.plugins.precision.mixed import MixedPrecisionPlugin -from pytorch_lightning.utilities import _NATIVE_AMP_AVAILABLE, _TORCH_GREATER_EQUAL_1_10, AMPType +from pytorch_lightning.utilities import _NATIVE_AMP_AVAILABLE, _TORCH_BFLOAT_AVAILABLE, AMPType from pytorch_lightning.utilities.exceptions import MisconfigurationException @@ -46,7 +46,7 @@ def __init__(self, precision: Union[int, str] = 16) -> None: def _select_precision_dtype(self, precision: Union[int, str] = 16) -> torch.dtype: if precision == "bf16": - if not _TORCH_GREATER_EQUAL_1_10: + if not _TORCH_BFLOAT_AVAILABLE: raise MisconfigurationException( "To use bfloat16 with native amp you must install torch greater or equal to 1.10." ) diff --git a/pytorch_lightning/utilities/__init__.py b/pytorch_lightning/utilities/__init__.py index 747c0be617cbd..56e3e03910de9 100644 --- a/pytorch_lightning/utilities/__init__.py +++ b/pytorch_lightning/utilities/__init__.py @@ -44,6 +44,7 @@ _OMEGACONF_AVAILABLE, _POPTORCH_AVAILABLE, _RICH_AVAILABLE, + _TORCH_BFLOAT_AVAILABLE, _TORCH_GREATER_EQUAL_1_7, _TORCH_GREATER_EQUAL_1_8, _TORCH_GREATER_EQUAL_1_9, diff --git a/pytorch_lightning/utilities/imports.py b/pytorch_lightning/utilities/imports.py index fa6598f884b19..dd31d23f6cf84 100644 --- a/pytorch_lightning/utilities/imports.py +++ b/pytorch_lightning/utilities/imports.py @@ -87,8 +87,13 @@ def _compare_version(package: str, op, version) -> bool: _OMEGACONF_AVAILABLE = _module_available("omegaconf") _POPTORCH_AVAILABLE = _module_available("poptorch") _RICH_AVAILABLE = _module_available("rich") +_TORCH_BFLOAT_AVAILABLE = _compare_version( + "torch", operator.ge, "1.10.0.dev20210820" +) # todo: swap to 1.10.0 once released _TORCH_QUANTIZE_AVAILABLE = bool([eg for eg in torch.backends.quantized.supported_engines if eg != "none"]) -_TORCH_SHARDED_TENSOR_AVAILABLE = _compare_version("torch", operator.ge, "1.10.0.dev20210809") +_TORCH_SHARDED_TENSOR_AVAILABLE = _compare_version( + "torch", operator.ge, "1.10.0.dev20210809" +) # todo: swap to 1.10.0 once released _TORCHTEXT_AVAILABLE = _module_available("torchtext") _TORCHVISION_AVAILABLE = _module_available("torchvision") _TORCHMETRICS_LOWER_THAN_0_3 = _compare_version("torchmetrics", operator.lt, "0.3.0") diff --git a/tests/models/test_amp.py b/tests/models/test_amp.py index 79c0cf7c12f15..173d1f8a0d1f8 100644 --- a/tests/models/test_amp.py +++ b/tests/models/test_amp.py @@ -22,7 +22,7 @@ import tests.helpers.utils as tutils from pytorch_lightning import Trainer from pytorch_lightning.plugins.environments import SLURMEnvironment -from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_10 +from pytorch_lightning.utilities import _TORCH_BFLOAT_AVAILABLE from pytorch_lightning.utilities.exceptions import MisconfigurationException from tests.helpers import BoringModel, RandomDataset from tests.helpers.runif import RunIf @@ -71,7 +71,7 @@ def predict(self, batch, batch_idx, dataloader_idx=None): 16, pytest.param( "bf16", - marks=pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_10, reason="torch.bfloat16 not available"), + marks=pytest.mark.skipif(not _TORCH_BFLOAT_AVAILABLE, reason="torch.bfloat16 not available"), ), ], )