src/transformers/models/flava/modeling_flava.py/FlavaGlobalContrastiveHead/forward
class FlavaGlobalContrastiveHead:    def forward(self, image_embeddings, text_embeddings, logit_scale):
        temperature = torch.exp(logit_scale)
        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
            labels = torch.arange(image_embeddings.size(0), device=image_embeddings.device)
            image_embeddings_all = [image_embeddings]
            text_embeddings_all = [text_embeddings]
        else:
            local_batch_size = image_embeddings.size(0)
            world_size = torch.distributed.get_world_size()

            if self.global_backprop_contrastive:
                image_embeddings_all = torch.distributed.nn.functional.all_gather_with_backprop(image_embeddings)
                text_embeddings_all = torch.distributed.nn.functional.all_gather_with_backprop(text_embeddings)
            else:
                image_embeddings_all = [torch.zeros_like(text_embeddings) for _ in range(world_size)]
                text_embeddings_all = [torch.zeros_like(image_embeddings) for _ in range(world_size)]
                torch.distributed.all_gather(image_embeddings_all, image_embeddings)
                torch.distributed.all_gather(text_embeddings_all, text_embeddings)

            labels = local_batch_size * torch.distributed.get_rank() + torch.arange(
                local_batch_size, device=image_embeddings.device
            )

        image_embeddings_all = torch.cat(image_embeddings_all)
        text_embeddings_all = torch.cat(text_embeddings_all)

        logits_per_image = torch.matmul(image_embeddings, text_embeddings_all.transpose(0, 1)) * temperature
        logits_per_text = torch.matmul(text_embeddings, image_embeddings_all.transpose(0, 1)) * temperature

        return logits_per_image, logits_per_text, labels