From 311bce3dac231d036921d1a661f6c0240afa993f Mon Sep 17 00:00:00 2001 From: SkafteNicki Date: Mon, 21 Aug 2023 11:53:32 +0200 Subject: [PATCH] fix mypy --- src/torchmetrics/utilities/imports.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/torchmetrics/utilities/imports.py b/src/torchmetrics/utilities/imports.py index 61c3e042fee..b7bc5969aca 100644 --- a/src/torchmetrics/utilities/imports.py +++ b/src/torchmetrics/utilities/imports.py @@ -30,10 +30,11 @@ _TORCH_GREATER_EQUAL_1_12: Optional[bool] = compare_version("torch", operator.ge, "1.12.0") _TORCH_GREATER_EQUAL_1_13: Optional[bool] = compare_version("torch", operator.ge, "1.13.0") -if module_available("lightning"): - _LIGHTNING_GREATER_EQUAL_2_0: Optional[bool] = compare_version("lightning", operator.ge, "2.0.0") -else: - _LIGHTNING_GREATER_EQUAL_2_0: Optional[bool] = compare_version("pytorch_lightning", operator.ge, "2.0.0") +_LIGHTNING_GREATER_EQUAL_2_0: Optional[bool] = ( + compare_version("lightning", operator.ge, "2.0.0") + if module_available("lightning") + else compare_version("pytorch_lightning", operator.ge, "2.0.0") +) _JIWER_AVAILABLE: bool = package_available("jiwer") _NLTK_AVAILABLE: bool = package_available("nltk")