Веса внимания вверху изображения

h = 16
fig, ax = plt.subplots(ncols=3, nrows=1, figsize=(15, 5))

for i, q_id in enumerate(sorted_indices[0]):
    logit = itm_logit[:, q_id, :]
    prob = torch.nn.functional.softmax(logit, dim=1)
    name = f'{prob[0, 1]:.3f}_query_id_{q_id}'
    
    # Attention map
    attention_map = avg_cross_att[0, q_id, :-1].view(h, h).detach().cpu().numpy()
    
    # Image
    raw_image_resized = raw_image.resize((596, 596))
    
    ax[0].set_title(name)
    ax[0].imshow(attention_map, cmap='viridis')
    ax[0].axis('off')
    
    ax[1].set_title(caption)
    ax[1].imshow(raw_image_resized)
    ax[1].axis('off')
    

    ax[2].set_title(f'Overlay: {name}')
    ax[2].imshow(raw_image_resized)
    ax[2].imshow(attention_map, cmap='viridis', alpha=0.6)  
    ax[2].axis('off')
    

    ax[0].set_aspect('equal')
    ax[1].set_aspect('equal')
    ax[2].set_aspect('equal')
    
    plt.tight_layout()
    plt.savefig(f"./att_maps/{name}.jpg")
    plt.show()
    break

Я пытаюсь наложить веса внимания поверх изображения (по треугольным осям), чтобы я мог видеть, на какой части веса внимания больше сосредоточено внимание.

Однако код, который я разместил, лишь перекрывает вес внимания поверх изображения.

В чем может быть проблема в этом случае?