Update resnet_lightning.py

resnet_lightning.py

@@ -16,7 +16,7 @@ from albumentations.pytorch import ToTensorV2
 
 
 from torch_lr_finder import LRFinder
-
+import lightning as L
 from pytorch_grad_cam import GradCAM
 from utils import *
 
@@ -45,7 +45,7 @@ class BasicBlock(nn.Module):
         return out
 
 
-class ResNet18Model(LightningModule):
+class ResNet18Model(L.LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, block=BasicBlock, num_blocks=[2, 2, 2, 2], num_classes=10):
         super(ResNet18Model, self).__init__()
         self.data_dir = data_dir