docs: def init_epoch(self):
          """Set up the batch generator for a new epoch."""
  
          if self._restored_from_state:
              self.random_shuffler.random_state = self._random_state_this_epoch
          else:
              self._random_state_this_epoch = self.random_shuffler.random_state
  
          self.create_batches()
  
          if self._restored_from_state:
              self._restored_from_state = False
          else:
              self._iterations_this_epoch = 0
  
          if not self.repeat:
              self.iterations = 0