Batched Cosine Similarity in PyTorch
Inputs:

Tensor a of shape [batch_size, n, d]
Tensor b of shape [batch_size, m, d]

Output:

Tensor c of shape [batch_size, n, m] where c[i, j, k] is the cosine similarity between a[i, j] and b[i, k]

How to implement this efficiently in PyTorch (preferably without for loops)?

try this:
c = torch.cosine_similarity(a.unsqueeze(2), b.unsqueeze(1), dim=-1)

