Flax's use_memory_efficient_attention is broken

[entrpn]

Describe the bug

I accidentally introduced a bug in this PR by making a condition on this line, which is needed when use_memory_efficient_attention=True.

Adding this bug to remind myself to put in a fix.

Reproduction

Using this example but adding use_memory_efficient_attention=True in from_pretrained.

Logs

No response

System Info

N/A

Who can help?

@pcuenca