ptrblck:

a = torch.randn(1, 2, 10, 10) b = torch.randn(1, 2, 10, 10) F.cosine_similarity(a.view(1, 2, -1), b.view(1, 2, -1), 2) > tensor([[-0.0755, 0.0896]])


@ptrblck: I am using toch.bmm to compute the cosine distance. it will return a maxtrix size of NxN instead of triangle vector in the matrix in the nn.CosineSimilarity. How to use nn.CosineSimilarity to get full cosine matrix as torch.bmm did? This is my code.
input1 = torch.randn(2, 4, 4)
input2 = torch.randn(2, 4, 4)
#Using bmm
x_norm = input1 / torch.norm(input1, p=2, dim=1, keepdim=True)
y_norm = input2 / torch.norm(input2, p=2, dim=1, keepdim=True)
cosine_sim = torch.bmm(x_norm.transpose(2,1), y_norm)
print('Using bmm: \n', cosine_sim)
# Pytorch built-in
cos = nn.CosineSimilarity(dim=1, eps=1e-6)
cosine_sim = cos(input1, input2)
print('Using nn: \n', cosine_sim)

The output is
Using bmm: 
 tensor([[[-0.0230,  0.2983,  0.0487,  0.3974],
         [-0.5747,  0.5513, -0.6436, -0.1389],
         [-0.3876, -0.2107,  0.7093, -0.4929],
         [-0.3446, -0.5347,  0.6372, -0.6423]],

        [[-0.3842, -0.0349,  0.1621,  0.6400],
         [ 0.6776, -0.4812, -0.3169, -0.7976],
         [-0.5251, -0.1258,  0.9381, -0.2379],
         [-0.1517,  0.7164,  0.8332,  0.1668]]])
Using nn: 
 tensor([[-0.0230,  0.5513,  0.7093, -0.6423],
        [-0.3842, -0.4812,  0.9381,  0.1668]])

Do you know how to use  nn.CosineSimilarity and achieve similar result as torch.bmm. I cannot use torch.bmm because of CUDA memory error.