def format_passages_batch(
        self,
        batch: MutableMapping,
        last_hidden_state: torch.Tensor,
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        """Format `passages` by selecting every ``n``th entry from the batch.

        Here ``n`` is the step size, which represents the number of hard
        negatives per passage.
        """
        assert self.step_size
        passages = {}
        num_blocks = batch['input_ids'].size(0) // self.step_size
        index = torch.arange(
            1,
            num_blocks * self.step_size + 1,
            device=last_hidden_state.device,
        ).view(num_blocks, self.step_size)
        index = index[:, :self.step_size - 1].reshape(-1)
        for key in batch:
            passages[key] = batch[key][index]
        return passages, last_hidden_state[index, :, :]
Format `passages` by selecting every ``n``th entry from the batch.

        Here ``n`` is the step size, which represents the number of hard
        negatives per passage.
        
llmfoundry/models/llm_embed/modeling_llm_embed.py