[Tune] Add support for new style lightning import (ray-project#44339)

Signed-off-by: marian-code <[email protected]>

python/ray/tune/integration/pytorch_lightning.py

@@ -6,7 +6,11 @@
 from contextlib import contextmanager
 from typing import Dict, List, Optional, Type, Union
 
-from pytorch_lightning import Callback, Trainer, LightningModule
+try:
+    from lightning import Callback, Trainer, LightningModule
+except ModuleNotFoundError:
+    from pytorch_lightning import Callback, Trainer, LightningModule
+
 from ray import train
 from ray.util import log_once
 from ray.util.annotations import PublicAPI, Deprecated