Fixed mypy errors in the backward compatibility code of LightningCLI …

…import
Lightning-Universe · Jul 2, 2023 · 7621b99 · 7621b99
1 parent 7c0b0b6
commit 7621b99
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/src/pl_bolts/models/detection/retinanet/retinanet_module.py b/src/pl_bolts/models/detection/retinanet/retinanet_module.py
@@ -135,14 +135,16 @@ def configure_optimizers(self):
 
 @under_review()
 def cli_main():
+ import pytorch_lightning
+
  try: # Backward compatibility for Lightning CLI
- from pytorch_lightning.cli import LightningCLI # PL v1.9+
- except ImportError:
- from pytorch_lightning.utilities.cli import LightningCLI # PL v1.8
+ cli_class: Any = getattr(pytorch_lightning.cli, "LightningCLI") # PL v1.9+
+ except AttributeError:
+ cli_class = getattr(pytorch_lightning.utilities.cli, "LightningCLI") # PL v1.8
 
  from pl_bolts.datamodules import VOCDetectionDataModule
 
- LightningCLI(RetinaNet, VOCDetectionDataModule, seed_everything_default=42)
+ cli_class(RetinaNet, VOCDetectionDataModule, seed_everything_default=42)
 
 
 if __name__ == "__main__":

diff --git a/src/pl_bolts/models/detection/yolo/yolo_module.py b/src/pl_bolts/models/detection/yolo/yolo_module.py
@@ -4,7 +4,6 @@
 import torch
 import torch.nn as nn
 from pytorch_lightning import LightningModule
-from pytorch_lightning.utilities.cli import LightningCLI
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 from torch import Tensor, optim
 
@@ -615,4 +614,11 @@ def _resize(self, image: Tensor, target: TARGET) -> Tuple[Tensor, TARGET]:
 
 
 if __name__ == "__main__":
- LightningCLI(CLIYOLO, ResizedVOCDetectionDataModule, seed_everything_default=42)
+ import pytorch_lightning
+
+ try: # Backward compatibility for Lightning CLI
+ cli_class: Any = getattr(pytorch_lightning.cli, "LightningCLI") # PL v1.9+
+ except AttributeError:
+ cli_class = getattr(pytorch_lightning.utilities.cli, "LightningCLI") # PL v1.8
+
+ cli_class(CLIYOLO, ResizedVOCDetectionDataModule, seed_everything_default=42)
diff --git a/src/pl_bolts/models/self_supervised/moco/moco_module.py b/src/pl_bolts/models/self_supervised/moco/moco_module.py
@@ -319,12 +319,14 @@ def _data_loader(self, dataset: Dataset, shuffle: bool = False) -> DataLoader:
 
 
 def cli_main() -> None:
+ import pytorch_lightning
+
  try: # Backward compatibility for Lightning CLI
- from pytorch_lightning.cli import LightningCLI # PL v1.9+
- except ImportError:
- from pytorch_lightning.utilities.cli import LightningCLI # PL v1.8
+ cli_class: Any = getattr(pytorch_lightning.cli, "LightningCLI") # PL v1.9+
+ except AttributeError:
+ cli_class = getattr(pytorch_lightning.utilities.cli, "LightningCLI") # PL v1.8
 
- LightningCLI(MoCo, CIFAR10ContrastiveDataModule, seed_everything_default=42)
+ cli_class(MoCo, CIFAR10ContrastiveDataModule, seed_everything_default=42)
 
 
 if __name__ == "__main__":