def METHOD_NAME(self, epoch, batch_size):
    self.indices = np.array([0] * len(self.buckets_idcs))
    self.step = 0
    self.epoch = epoch
    self.batch_size = batch_size
    self.steps_per_epoch = self._compute_steps_per_epoch()