Update display.py

display.py

@@ -9,7 +9,7 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 from utils import YoloCAM, cells_to_bboxes, non_max_suppression
 
 def inference(image: np.ndarray, iou_thresh: float = 0.75, thresh: float = 0.75,show_cam: bool = False,  transparency: float = 0.5):
-    model = YOLOV3_PL()
+    model = YOLOv3(num_classes=20)
     model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
     # iou_thresh = 0.75
     # thresh = 0.75