Currently, shuffle does not work properly, which is good enough for a lot of use cases, but can probably result in lower accuracy when there are many training epochs. This should be resolved by properly permuting the dataset across ranks.

In addition, there needs to be a check that raises a warning if the number of batches on a rank is 0, since this is almost certainly going to cause training to hang (which is normal PyTorch behavior, but the user should be warned first).