Unet model only supports internal block channel sizes in multiples of 32

[ezhang7423]

What API design would you like to have changed or added to the library? Why?

I see that the current unet model breaks if block_out_channels is set to something other than a multiple of 32. I am just curious if this was on purpose? I would actually like to use a size smaller than 32. I tested this by changing the first element in the tuple on line 54 from 32 to 31 intests/test_models_unet.py and observe some of the tests fail afterwards.

What use case would this enable or better enable? Can you give us a code example?

When debugging or playing around with small models, it would be nice to have the capability to set the batch_out_channels to be smaller than 32.

I would be happy to work on supporting this feature if it seems useful.

[patrickvonplaten]

Agree! Happy to adapt the code to make it work for numbers different to 32!

Could you maybe copy-post a short reproducible code snippet that you would like to have functioning? :-)
This would really help to better understand your use case.

[patil-suraj]

Hey @ezhang7423 !Thank you for opening the issue.

It needs to be multiple of 32 because the norm_num_groups is set to 32. Opened a PR to fix this.
Also note that the block_out_channels values needs to be multiple of norm_num_groups.