def METHOD_NAME(
    self,
    id_list: List[str],
    batches: Dict[str, List[torch.Tensor]],
    shuffle: bool,
    state: np.random.RandomState,
):
    if shuffle:
        indices = np.arange(0, len(id_list))
        state.shuffle(indices)
        batches = {k: [v[i] for i in indices] for k, v in batches.items()}
        id_list = [id_list[i] for i in indices]
    bs = self.batch_size
    while len(id_list) >= bs:
        # Make mini-batch and yield
        yield (
            id_list[:bs],
            {k: torch.stack(v[:bs], 0) for k, v in batches.items()},
        )
        id_list = id_list[bs:]
        batches = {k: v[bs:] for k, v in batches.items()}
    return id_list, batches