masked_simple = attn_weights*mask_simple
print(masked_simple)