Thank you very much for your reply!
If mat is the predicted probability matrix with shape (128,5), then torch.mm(mat,mat.t()) gives me the 128 by 128 matrix that contains all the inner products between pairs of rows p,q in mat.
But what I’m hoping is to compute a more general function k(p,q) between all pairs of rows p,q in mat and store it in a 128 by 128 matrix. So torch.mm(mat,mat.t()) can be seen as a simple case of this where the function k(p,q) is just the inner product <p,q>.