Robust layer extraction for GradCAM

xai_engine/methods.py

@@ -67,10 +67,14 @@ 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."""
-    layer = wrapper.model.vit.encoder.layer[target_layer_idx]
+    model = wrapper.model if hasattr(wrapper, "model") else wrapper
+    vit_base = model.vit if hasattr(model, "vit") else model
+    encoder = vit_base.encoder if hasattr(vit_base, "encoder") else vit_base
+    layer = 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)
         result = tensor[:, 1:, :].reshape(tensor.size(0), 24, 24, tensor.size(2))