Try lightning when pytorch_lightning is not found (#10404)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: akihironitta

CHANGELOG.md

@@ -7,8 +7,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- Fix `detach()` warnings in example scripts involving tensor conversions. ([#10357](https://github.com/pyg-team/pytorch_geometric/pull/10357))
-- Fix non-tuple indexing to resolve PyTorch deprecation warning. ([#10389](https://github.com/pyg-team/pytorch_geometric/pull/10389))
+- Fixed importing PyTorch Lightning in `torch_geometric.graphgym` and `torch_geometric.data.lightning` when using `lightning` instead of `pytorch-lightning` ([#10404](https://github.com/pyg-team/pytorch_geometric/pull/10404))
+- Fixed `detach()` warnings in example scripts involving tensor conversions ([#10357](https://github.com/pyg-team/pytorch_geometric/pull/10357))
+- Fixed non-tuple indexing to resolve PyTorch deprecation warning ([#10389](https://github.com/pyg-team/pytorch_geometric/pull/10389))
 
 ### Added
 

torch_geometric/data/lightning/datamodule.py

@@ -11,21 +11,27 @@
 from torch_geometric.typing import InputEdges, InputNodes, OptTensor
 
 try:
-    from pytorch_lightning import LightningDataModule as PLLightningDataModule
-    no_pytorch_lightning = False
+    from pytorch_lightning import LightningDataModule as _LightningDataModule
+    _pl_is_available = True
 except ImportError:
-    PLLightningDataModule = object  # type: ignore
-    no_pytorch_lightning = True
+    try:
+        from lightning.pytorch import \
+            LightningDataModule as _LightningDataModule
+        _pl_is_available = True
+    except ImportError:
+        _pl_is_available = False
+        _LightningDataModule = object
 
 
-class LightningDataModule(PLLightningDataModule):
+class LightningDataModule(_LightningDataModule):
     def __init__(self, has_val: bool, has_test: bool, **kwargs: Any) -> None:
         super().__init__()
 
-        if no_pytorch_lightning:
+        if not _pl_is_available:
             raise ModuleNotFoundError(
-                "No module named 'pytorch_lightning' found on this machine. "
-                "Run 'pip install pytorch_lightning' to install the library.")
+                "No module named 'pytorch_lightning' (or 'lightning') found "
+                "in your Python environment. Run 'pip install "
+                "pytorch_lightning' or 'pip install lightning'")
 
         if not has_val:
             self.val_dataloader = None  # type: ignore

torch_geometric/graphgym/imports.py

@@ -4,14 +4,23 @@
 
 try:
     import pytorch_lightning as pl
+    _pl_is_available = True
+except ImportError:
+    try:
+        import lightning.pytorch as pl
+        _pl_is_available = True
+    except ImportError:
+        _pl_is_available = False
+
+if _pl_is_available:
     LightningModule = pl.LightningModule
     Callback = pl.Callback
-except ImportError:
+else:
     pl = object
     LightningModule = torch.nn.Module
     Callback = object
 
     warnings.warn(
-        "Please install 'pytorch_lightning' via  "
-        "'pip install pytorch_lightning' in order to use GraphGym",
+        "To use GraphGym, install 'pytorch_lightning' or 'lightning' via "
+        "'pip install pytorch_lightning' or 'pip install lightning'",
         stacklevel=2)