torch row wise cosinus similarity
Assuming I am having 2 tensors with dimension:

tensor A: [512, 100]
tensor B: [512, 200, 100]

and I would like to use torch.cosine_similarity to produce a tensor C of dimensions [512, 1, 200] or [512, 200].
Using torch.cosine_similarity(A, B) I get an error:
'{RuntimeError}The size of tensor a (512) must match the size of tensor b (200) at non-singleton dimension 1'
I guess it can be done by the following:
desired_result = torch.stack([torch.cosine_similarity(A_row, B_row, axis=-1) for B_row, A_row in zip(B, A)])

but there should be a more optimized way. Any help/hint?

No need for stacking, broadcasting will do the job for you:
# (512, 100, 1) after unsqueeze
A = torch.randn(512, 100).unsqueeze(dim=-1)
B = torch.randn(512, 200, 100)

# (512, 200)
torch.cosine_similarity(B, A.unsqueeze(1), axis=-1)

