Visualizing the attention map of a multihead attention in ViT

I'm trying to visualize the attention map of mit Visual Transformer architecture in keras/tensorflow. For this I was able to implement the ViT model the following way:

def model():
    input_layer = layers.Input(shape=input_shape)
    #image_patches = create_patches(input_layer)
    #print(input_layer.shape)
    image_patches = Patches(patch_size)(input_layer)
    #print(image_patches.shape)
    encoded_patches = PatchEncoder(num_patch, projection_dim)(image_patches)
    #print(encoded_patches.shape)
    
    #for i in range(transformer_blocks):
    x1 = layers.LayerNormalization()(encoded_patches)
    x1 = layers.MultiHeadAttention(num_heads=num_heads, key_dim=projection_dim, name='MHA_1')(x1, x1)
    x = layers.Add()([x1, encoded_patches])
    x2 = layers.LayerNormalization()(x)
    x2 = mlp_head(x2, transformer_units)
    encoded_patches = layers.Add()([x2, x])    
    
    x = layers.LayerNormalization()(encoded_patches)
    x = layers.Flatten()(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(2)(x)
    
    model = tf.keras.Model(inputs=input_layer, outputs=x)
    print(model.summary())
    return model    

I'm now trying to visualize the attention map based on an input image and my model output. For this I first try to predict the outcome and reshape the weights:

def attention_map(model, image):
    size = model.input_shape[1]
    grid_size = int(np.sqrt(model.layers[4].output_shape[-2] - 1))
    
    # Prepare the input
    X = preprocess_inputs(cv2.resize(image, (size, size)))#[np.newaxis, :]  # type: ignore

    # Get the attention weights from each transformer.
    outputs = [
        l.output[1] for l in model.layers if isinstance(l, layers.MultiHeadAttention)
    ]
    weights = np.array(
        tf.keras.models.Model(inputs=model.inputs, outputs=outputs).predict(X_test)
    )
    print(weights.shape)
    num_layers = weights.shape[0]
    num_heads = weights.shape[1]
    reshaped = weights.reshape(
        (num_layers, num_heads, grid_size ** 2 + 1, grid_size ** 2 + 1)
    )

    # From Appendix D.6 in the paper ...
    # Average the attention weights across all heads.
    reshaped = reshaped.mean(axis=1)

    # From Section 3 in https://arxiv.org/pdf/2005.00928.pdf ...
    # To account for residual connections, we add an identity matrix to the
    # attention matrix and re-normalize the weights.
    reshaped = reshaped + np.eye(reshaped.shape[1])
    reshaped = reshaped / reshaped.sum(axis=(1, 2))[:, np.newaxis, np.newaxis]

    # Recursively multiply the weight matrices
    v = reshaped[-1]
    for n in range(1, len(reshaped)):
        v = np.matmul(v, reshaped[-1 - n])

    # Attention from the output token to the input space.
    mask = v[0, 1:].reshape(grid_size, grid_size)
    mask = cv2.resize(mask / mask.max(), (image.shape[1], image.shape[0]))[
        ..., np.newaxis
    ]
    return (mask * image).astype("uint8")

However my problem is now to reshape my weight matrix getting in mismatch. Can someone give me a hint on why this is occuring? A hint based on the output dimension given by

    weights = np.array(
        tf.keras.models.Model(inputs=model.inputs, outputs=outputs).predict(X_test)
    )

would also help.