def _cat_gather(self, t: torch.Tensor, group: Any = None) -> torch.Tensor:
        """Applies an all gather operation necessary for InfoNCELoss.

        See
        https://github.com/pytorch/pytorch/blob/63d5e9221bedd1546b7d364b5ce4171547db12a9/torch/distributed/nn/functional.py#L314
         as well as
        https://github.com/pytorch/pytorch/issues/121587#issuecomment-1989070351
        """
        if self.use_legacy_gradient_passthrough:
            all_tensors = list(dist.all_gather(t, group))
            all_tensors[dist.get_global_rank()] = t
            all_tensors = torch.cat(all_tensors)
        else:
            extra_kwargs = {'group': group} if group is not None else {}
            all_tensors = all_gather(t, **extra_kwargs)
            all_tensors = torch.cat(all_tensors)

        return all_tensors
Applies an all gather operation necessary for InfoNCELoss.

        See
        https://github.com/pytorch/pytorch/blob/63d5e9221bedd1546b7d364b5ce4171547db12a9/torch/distributed/nn/functional.py#L314
         as well as
        https://github.com/pytorch/pytorch/issues/121587#issuecomment-1989070351
        
llmfoundry/models/llm_embed/modeling_llm_embed.py