Well, if it’s of any use, I’ve posted my code below. Since this doesn’t use any fancy torch or even numpy stuff, it’s almost trivial. In my case, inputs is a list of N encoded sequences with a maximum length max_seq_len that are padded if needed. lengths is a list with lengths for each of the N sequences (max_seq_len - padding) – I have my own vectorizer which returns lengths with the encoded inputs.
With this I can simply call
train_iter = BucketDataset(X_train, y_train, lengths_train, batch_size)

with my training loop looks like this:
for idx, (inputs, lengths) in enumerate(train_iter):
    batch_size = inputs.shape[0]
    # Remove unnecessary padding (all sequences have the same length)
    inputs = inputs[:, 0:lengths[0]]
    # Convert to tensors and move to device
    inputs = torch.tensor(inputs).to(self.device)
    ...

Uh, that’s it :). Of course, some batches might be smaller than batch_size, but if the dataset is reasonably large those number is almost negligible. For example, when I use a dataset of mine with about 840k sentences, 99.7% of batches are full and the rest are almost full. So there’s no measurable performance loss here.
from collections import OrderedDict

class BucketDataset:

    def __init__(self, inputs, targets, lengths, batch_size):
        self.inputs = inputs      # shape = (N, max_seq_len)
        self.targets = targets    # shape = (N, ) or None (e.g., for autoencoder I can simply use inputs)
        self.lengths = lengths    # shape = (N, )
        self.batch_size = batch_size
        self.current = -1
        self._generate_batch_map()

    def _generate_batch_map(self, equal_length=False):
        batch_map = OrderedDict()
        # Organize lengths, e.g., batch_map[10] = [30, 124, 203, ...] <= indices of sequences of length 10
        for idx, length in enumerate(self.lengths):
            if length not in batch_map:
                batch_map[length] = [idx]
            else:
                batch_map[length].append(idx)
        # Use batch_map to split indices into batches of equal size
        # e.g., for batch_size=3, batch_list = [[23,45,47], [49,50,62], [63,65,66], ...]
        self.batch_list = []
        for length, indices in batch_map.items():
            for group in [indices[i:(i+self.batch_size)] for i in range(0, len(indices), self.batch_size)]:
                self.batch_list.append(group)

    def batch_count(self):
        return len(self.batch_list)

    def __len__(self):
        return len(self.lengths)

    def __iter__(self):
        return self

    def __next__(self):
        self.current += 1
        if self.current > len(self.batch_list)-1:
            self.current = -1
            raise StopIteration
        else:
            if self.targets is None:
                return self.inputs[self.batch_list[self.current]], \
                       self.lengths[self.batch_list[self.current]]
            else:
                return self.inputs[self.batch_list[self.current]], \
                       self.targets[self.batch_list[self.current]], \
                       self.lengths[self.batch_list[self.current]]