I’m trying to implement VIT Attention Rollout from github Vit Explain using PyTorch’s Vision Transformer model, but I’m encountering an IndexError: list index out of range error. This error seems to be related to the fused_attn attribute in the attention mechanism.

Here’s the relevant part of my code:

def load_model(weight_path, device):
    model = torchvision.models.vit_l_16()
    model.heads = torch.nn.Linear(
        in_features=model.heads.head.in_features, out_features=2)
    model.load_state_dict(torch.load(
        weight_path, map_location=torch.device(device)))
    model.eval()
    return model

model = load_model(weight_path=path, device=device)
img = preprocess_image(image_path)
attention_rollout = VITAttentionRollout(model, head_fusion="max", discard_ratio=0.9)
mask = attention_rollout(img)

Error Message

IndexError                                Traceback (most recent call last)
Cell In[48], line 2
      1 attention_rollout = VITAttentionRollout(model, head_fusion="max", discard_ratio=0.9)
----> 2 mask = attention_rollout(img)

Cell In[46], line 130
    127 with torch.no_grad():
    128     output = self.model(input_tensor)
--> 130 return rollout(self.attentions, self.discard_ratio, self.head_fusion)

Cell In[46], line 75
     74 def rollout(attentions, discard_ratio, head_fusion):
---> 75     result = torch.eye(attentions[0].size(-1))
     76     with torch.no_grad():
     77         for attention in attentions:

IndexError: list index out of range

I found a similar issue with the vit_explain library from GitHub, where the problem was resolved by setting attn.fused_attn = False for each block in the model. However, I’m using PyTorch’s built-in Vision Transformer model, which doesn’t have this attribute.

the question is

1. How can I implement a similar fix (setting fused_attn = False) for PyTorch’s Vision Transformer model?
2. Is there a way to convert the PyTorch Vision Transformer architecture to use the timm library, which has the fused_attn attribute?
3. Are there any other approaches to resolve this IndexError when implementing VIT Attention Rollout with PyTorch’s Vision Transformer?

Any insights or suggestions would be greatly appreciated. Thank you!