I had the same problem when I re-implement PFE 20, which needs 64 classes, which have 4  data in each class, as an input.
Here’s my solution (not very good but useful for me).
#class Dataset
def __init__(self):
    self.label_to_index = {'A': [1, 3, 4], 'B': [0, 2, 5], 'C': [6, 7, 8]}

def __getitem__(self, item):
    images = np.empty((1, *image_shape))
    labels = np.empty((1,))
    for i, (label, indices) in enumerate(self.label_to_index.items()):
        index = np.random.choice(indices, 1, replacement=True)[0]
        images[i, ...] = get_image(index)
        labels[i] = get_label(index)
    return torch.from_numpy(images), torch.from_numpy(labels)

As you can see, it returns (1 from A, 1 from B, 1 from C) as an item. Then I set batch_size = 10.
I think, if the dataset is large enough, after a large number of iterations, random choice in each batch would be similar to shuffle the dataset then  iterate over it