src/lighteval/models/nanotron/nanotron_model.py [670:691]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            batch_size = self._get_batch_size(
                max_input_length=max_context,
                override_bs=self.batch_size,
                starting_batch_size=starting_batch_size,
            )

            starting_batch_size = batch_size * 2  # for the next round

            # For the DP replicas
            distributed_sampler = DistributedSampler(
                dataset,
                num_replicas=self.parallel_context.dp_pg.size(),
                rank=dist.get_rank(self.parallel_context.dp_pg),
                shuffle=False,
                drop_last=False,
            )
            to_remove_at_the_end = distributed_sampler.total_size - len(dataset)

            dataloader = DataLoader(
                dataset,
                batch_size=batch_size,
                sampler=distributed_sampler,
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



src/lighteval/models/nanotron/nanotron_model.py [908:928]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            batch_size = self._get_batch_size(
                max_input_length=max_context,
                override_bs=self.batch_size,
                starting_batch_size=starting_batch_size,
            )
            starting_batch_size = batch_size * 2  # for the next round

            # For the DP replicas
            distributed_sampler = DistributedSampler(
                dataset,
                num_replicas=self.parallel_context.dp_pg.size(),
                rank=dist.get_rank(self.parallel_context.dp_pg),
                shuffle=False,
                drop_last=False,
            )
            to_remove_at_the_end = distributed_sampler.total_size - len(dataset)

            dataloader = DataLoader(
                dataset,
                batch_size=batch_size,
                sampler=distributed_sampler,
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -