From what I could understand, nn.CosineSimilarity loss computes the cosine similarity between an element i of batch u and another element i of batch v. What I’m looking for is an approach to compute the similarity matrix of all elements of u to all elements of v and define it as a PyTorch loss function. It was so easy in Tensorflow 2.