def __init__(self, total_size, group_size):
        """
        Args:
            size (int): the total number of data of the underlying dataset to sample from
        """
        assert total_size > 0 and group_size > 0
        assert total_size % group_size == 0, \
           f"The total size must be divisible by group size: total size={total_size}, group size={group_size}"

        self._total_size = total_size
        self._group_size = group_size
        self._rank = comm.get_rank()
        self._world_size = comm.get_world_size()

        self._num_groups = total_size // group_size

        shard_size = ((self._num_groups - 1) // self._world_size + 1) * self._group_size

        # shard_size = (self._total_size - 1) // self._world_size + 1
        begin = shard_size * self._rank
        end = min(shard_size * (self._rank + 1), self._total_size)
        self._local_indices = range(begin, end)

        Args:
            size (int): the total number of data of the underlying dataset to sample from
        
tridet/data/samplers/group_sampler.py