    def test_CoCa(self):
        """
        Test CoCa model.
        """
        model = nn.Sequential()
        coca = CoCa(
            dim = 4,
            img_encoder = model,
            image_dim = 32,
            num_tokens = 10,
            unimodal_depth = 6,
            multimodal_depth = 6,
            dim_head = 64,
            heads = 8,
            caption_loss_weight = 1.,
            contrastive_loss_weight = 1.,
        )

        text = torch.randint(0, 10, (4, 4))
        images = torch.randn(1, 2, 32)
        logits = coca(
            text = text,
            images = images
        )
        self.assertTrue(logits.shape == torch.Size([4,4,10]))
        text_embeds, image_embeds = coca(
            text = text,
            images = images,
            return_embeddings = True
        )
        self.assertTrue(text_embeds.shape == torch.Size([4,4]))
        self.assertTrue(image_embeds.shape == torch.Size([1,4]))
Test CoCa model.
tests/unittests/models/coca/test_coca.py