def __iter__(self):
    if self.shuffle:
        batch = [[(i + (j * self.num_inst)) for i in torch.randperm(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
    else:
        batch = [[(i + (j * self.num_inst)) for i in range(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
    batch = [item for sublist in batch for item in sublist]
    if self.shuffle:
        random.shuffle(batch)
    return iter(batch)
def __iter__(self):
    if self.shuffle:
        batch = [[(i + (j * self.num_inst)) for i in torch.randperm(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
    else:
        batch = [[(i + (j * self.num_inst)) for i in range(self.num_inst)[:self.num_per_class]] for j in range(self.num_cl)]
    batch = [item for sublist in batch for item in sublist]
    if self.shuffle:
        random.shuffle(batch)
    return iter(batch)