Feature Request: Custom Buckets

[Vega-KH]

When doing aspect ratio bucketing, there are several predefined bucket groups, defined in dataloaders_util.py. For example, if I choose the 640 bucket group, I will get buckets like 640x640, 704x576, 768x512, 896x448, etc.

For my experiments, I would like to use a bucket that ALWAYS has 512 as one of dimensions. So the buckets would be like 512x512, 512x576, 512x640, 512x768, etc. I realize that the different buckets would have different memory requirements, but there are some reasons why this might work better.

To solve this myself, I tried to just edit the buckets in the dataloaders_util.py myself, but I get an immediate error after doing so, and I'm not familiar enough with the codebase to figure out the reason. Please consider making a way to add custom bucket groups.