fix mypy

Lightning-AI · Aug 21, 2023 · 311bce3 · 311bce3
1 parent 26ed08d
commit 311bce3
Showing 1 changed file with 5 additions and 4 deletions.
diff --git 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")