def _update_step_size_if_needed(self, batch: MutableMapping) -> None:
        """Update step size on first batch if we detect hard negatives."""
        if self.step_size:
            return

        input_shape = batch['input_ids'].shape
        if input_shape[1] > 2:
            # We have hard negatives, batch shape is [batch, sample of query+positive passage+negative passages, tokens].
            self.step_size = input_shape[1]
            log.info(
                f'Detected hard negatives, updated step_size to {self.step_size}',
            )
        else:
            self.step_size = 2
Update step size on first batch if we detect hard negatives.
llmfoundry/models/llm_embed/modeling_llm_embed.py