Fix GradCAM target layer reference for transformers ViT model

xai_engine/methods.py

@@ -67,8 +67,9 @@ def integrated_gradients(wrapper, input_tensor, steps=20) -> tuple[np.ndarray, f
 
 def vit_grad_cam(wrapper, input_tensor, target_layer_idx=11) -> np.ndarray:
     """GradCAM for ViT using pytorch_grad_cam with custom reshape transform."""
-    # Target layer: last encoder layer output
-    target_layer = [wrapper.model.vit.encoder.layer[target_layer_idx].output]
+    layer = wrapper.model.vit.encoder.layer[target_layer_idx]
+    target_layer = [layer.output if hasattr(layer, "output") else layer]
+
 
     def reshape_transform(tensor):
         # tensor shape: (B, 577, 384) -> drop CLS token -> (B, 576, 384) -> (B, 24, 24, 384) -> (B, 384, 24, 24)