contrastive_labels - dist batch size isn't correct

[gillotte]

CoCa-pytorch/coca_pytorch/coca_pytorch.py

Line 532 in edee92c

contrastive_labels = torch.arange(batch, device=device)

fyi the dist batch size isn't correct

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