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

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.