search

lucidrains / CoCa-pytorch

Implementation of CoCa, Contrastive Captioners are Image-Text Foundation Models, in Pytorch
MIT License
1.04k stars 88 forks source link

contrastive_labels - dist batch size isn't correct #20

Open gillotte opened 5 months ago

gillotte commented 5 months ago

https://github.com/lucidrains/CoCa-pytorch/blob/edee92c74e311ccfa4a0024412fd991c98aff5fd/coca_pytorch/coca_pytorch.py#L532

fyi the dist batch size isn't correct

torch.arange(batch, device=device) -> torch.arange(text_latents.shape[0], device=device)