Update model.py

model.py

@@ -1,7 +1,7 @@
 """
 Implementation of YOLOv3 architecture
 """
-
+from pytorch_lightning import LightningModule
 import torch
 import torch.nn as nn
 
@@ -100,7 +100,7 @@ class ScalePrediction(nn.Module):
         )
 
 
-class YOLOv3(nn.Module):
+class YOLOv3(LightningModule):
     def __init__(self, in_channels=3, num_classes=80):
         super().__init__()
         self.num_classes = num_classes