The docs give you an example:

input1 = torch.randn(100, 128)
input2 = torch.randn(100, 128)
output = F.cosine_similarity(input1, input2)
print(output)
If you want to use more dimensions, refer to the docs for the shape explanation.

E.g. for a 4-dim tensor, where you would like to compute the distance along dim2, this code should work:
input1 = torch.randn(100, 128, 32, 32)
input2 = torch.randn(100, 128, 32, 32)
output = F.cosine_similarity(input1, input2, dim=2)
print(output)
              
I give you my current example.

I obtain two tensor which are [1,3,512,512] by neural network.

I mean that I want to calculate cosine similarity of four-dim tensor.

Can I use torch.cosine_similarity(input1,input2,dim=3)?

Or I only compute the sum of each dimension first, and then average each dimension.
 JasonChenhx:
Or I only compute the sum of each dimension first, and then average each dimension.
That might depend on your use case and what you would like to achieve, so I can’t be really helpful here.