How to calculate cosine similarity of two multi-demensional vectors?

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)