Update display.py

display.py

@@ -15,7 +15,11 @@ def inference(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.5,sh
     # iou_thresh = 0.75
     # thresh = 0.75
     scaled_anchors  =  config.SCALED_ANCHORS
-    cam = YoloCAM(model=model, target_layers=[model.layers.26.conv,model.layers.28.conv], use_cuda=False)
+    
+    backbone = model
+    target_layer_list = list(backbone.children())[-2:]
+
+    cam = YoloCAM(model=model, target_layers = target_layer_list, use_cuda=False)
     transforms =  A.Compose(
     [
         A.LongestMaxSize(max_size=config.IMAGE_SIZE),