def on_epoch_end(self):
    if self.shuffle_groups:
        random.shuffle(self.groups)
    self.current_index = 0
def on_epoch_end(self):
    if self.shuffle_groups:
        random.shuffle(self.groups)
    self.current_index = 0