If you’re using .data, you’re almost certainly doing it wrong with tracing (and it’s not too great an idea in general).
If I’m understanding this correctly, you want the mask to leave data intact by what is not masked and set attn to -math.inf where it isn’t. The easiest way might be to take the 1, 0-valued (float) mask and do attn = attn + mask.log().
Best regards
Thomas