diffusers

a5f35ee4 - add reshape to fix use_memory_efficient_attention in flax (#7918)

add reshape to fix use_memory_efficient_attention in flax (#7918) Co-authored-by: Juan Acevedo <jfacevedo@google.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: Aryan <aryan@huggingface.co>