ZEYN

I am actually facing the same issue could you solve this problem ?

Actually I have deleted this line (it is not the best solution maybe) and had no problem running the code. Hopefully we can get the right solution

`# for attn in attn_weights: # # attn_weights_mean.append(attn[..., skip_len:(skip_len+k_w**2)].mean(dim=1).squeeze().cpu()) # attn_weights_mean.append(attn[..., skip_len:(skip_len+k_w**2)].mean(dim=1).squeeze().reshape(shape1+shape2).cpu())` just removed those lines