Update display.py

display.py

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