class_docstr: Sampler that partitions dataset across ranks and repartitions after reset events.
  
  Works similar to `DistributedSampler`, but with an optional capability to record
  which dataset indices have been processed each batch. When tracked by a `TorchState`
  object, the sampler will automatically repartition the unprocessed indices among the
  new set of workers.
  
  In order to use this object successfully it is recommended that the user:
  
  1. Include this object in the `TorchState`.
  2. Call `record_batch` after processing a set of samples.
  3. Call `set_epoch` at the end of each epoch to clear the processed indices.
  
  Args:
      dataset: Dataset used for sampling (assumed to be of constant size).
      shuffle: If `True` (default), shuffle the indices.
      seed: Random seed used to shuffle the sampler when `shuffle=True`.
            This number should be identical across all ranks (default: 0).
class_docstr_tok: ["Sampler","that","partitions","dataset","across","ranks","and","repartitions","after","reset","events",".","Works","similar","to","`","DistributedSampler","`",",","but","with","an","optional","capability","to","record","which","dataset","indices","have","been","processed","each","batch",".","When","tracked","by","a","`","TorchState","`","object",",","the","sampler","will","automatically","repartition","the","unprocessed","indices","among","the","new","set","of","workers",".","In","order","to","use","this","object","successfully","it","is","recommended","that","the","user",":","1",".","Include","this","object","in","the","`","TorchState","`",".","2",".","Call","`","record_batch","`","after","processing","a","set","of","samples",".","3",".","Call","`","set_epoch","`","at","the","end","of","each","epoch","to","clear","the","processed","indices",".","Args",":","dataset",":","Dataset","used","for","sampling","(","assumed","to","be","of","constant","size",")",".","shuffle",":","If","`","True","`","(","default",")",",","shuffle","the","indices",".","seed",":","Random","seed","used","to","shuffle","the","sampler","when","`","shuffle=True","`",".","This","number","should","be","identical","across","all","ranks","(","default",":","0",")","."]
docstr: Sets the epoch for this sampler.
  
  When `shuffle=True`, this ensures all replicas use a different random ordering
  for each epoch.
  
  Will clear and reset the `processed_indices` for the next epoch. It is important
  that this is called at the end of the epoch (not the beginning) to ensure that
  partially completed epochs do not reprocess samples.
  
  Args:
      epoch: Epoch number.
docstr_tok: ["Sets","the","epoch","for","this","sampler",".","When","`","shuffle=True","`",",","this","ensures","all","replicas","use","a","different","random","ordering","for","each","epoch",".","Will","clear","and","reset","the","`","processed_indices","`","for","the","next","epoch",".","It","is","important","that","this","is","called","at","the","end","of","the","epoch","(","not","the","beginning",")","to","ensure","that","partially","completed","epochs","do","not","reprocess","samples",".","Args",":","epoch",":","Epoch","number","."]
code: def set_epoch(self, epoch):
      """Sets the epoch for this sampler.
  
      When `shuffle=True`, this ensures all replicas use a different random ordering
      for each epoch.
  
      Will clear and reset the `processed_indices` for the next epoch. It is important
      that this is called at the end of the epoch (not the beginning) to ensure that
      partially completed epochs do not reprocess samples.
  
      Args:
          epoch: Epoch number.
      """
      self.epoch = epoch
      self.processed_num = 0
      self.reset()
  
file_path: horovod/horovod/torch/elastic/sampler.py
imports: import math
  import random
  import torch.utils.data.distributed
  from horovod.torch.mpi_ops import rank, size
label_desc: Use image node_id 2 for calling the ElasticSampler obj's underlying member method code with example usage: obj.set_epoch(epoch) without return types