Update import 'pytorch_lightning' -> 'lightning.pytorch' (#7)

* change import

Signed-off-by: Maanu Grover <maanug@nvidia.com>

* isort

Signed-off-by: Maanu Grover <maanug@nvidia.com>

---------

Signed-off-by: Maanu Grover <maanug@nvidia.com>

Author: maanug-nv

src/nvidia_resiliency_ext/ptl_resiliency/fault_tolerance_callback.py

@@ -26,7 +26,7 @@
 from typing import Optional, Union
 
 import torch
-from pytorch_lightning.callbacks import Callback
+from lightning.pytorch.callbacks import Callback
 
 import nvidia_resiliency_ext.fault_tolerance as ft
 

src/nvidia_resiliency_ext/ptl_resiliency/straggler_det_callback.py

@@ -19,7 +19,7 @@
 from typing import Optional
 
 import torch
-from pytorch_lightning.callbacks import Callback
+from lightning.pytorch.callbacks import Callback
 
 import nvidia_resiliency_ext.straggler as straggler
 

tests/ptl_resiliency/unit/test_ft_callback.py

@@ -24,11 +24,11 @@
 import sys
 import tempfile
 
+import lightning.pytorch as pl
 import pytest
-import pytorch_lightning as pl
 import torch
-from pytorch_lightning.callbacks import Callback
-from pytorch_lightning.utilities.exceptions import _TunerExitException
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.utilities.exceptions import _TunerExitException
 from torch import nn
 
 import nvidia_resiliency_ext.fault_tolerance as fault_tolerance

tests/ptl_resiliency/unit/test_straggler_det_callback.py

@@ -19,8 +19,8 @@
 import shutil
 import tempfile
 
+import lightning.pytorch as pl
 import pytest
-import pytorch_lightning as pl
 import torch
 from torch import nn