[Community] Add dropout to Flax UNet

[patrickvonplaten]

Describe the bug

Flax UNet currently doesn't support dropout

Reproduction

No response

Logs

No response

System Info

Main

[MarkRich]

Hello @patrickvonplaten !

I assume this issue is still relevant? If so happy to help with it. Just for clarity: the task is to address the comments that patil had here: #502 (review), correct? i.e. that the flax UNet model does not actually perform dropout since it was originally intended just for inference?

[patrickvonplaten]

Hey @MarkRich,

That's exactly right. We simply need to add dropout keywords to a lot of places so that it can be customized by the user

[patrickvonplaten]

I won't have time to look into this anytime soon, so unassigning from me.