@fmassa Thanks for your solution. I definitely need to refresh my linear algebra skills. It really solves the first part.
About, if you are not backpropagating through y part…I am a bit confused. Essentially, the algorithm is:

Build a CNN that on the final layer has 2 neurons.
Transform the output of those 2 neurons (a tensor of shape n x 2) to similarity matrix. Call it X.
Transform the labels y into a n x n tensor (where ij-th element is 1 if i and j belong to the same cluster, 0 otherwise). Call it Y.
Do an elementwise multiplication of X and Y.
Do the extra stuff, sum, some substraction etc.

While I do not need to backprop through y, Y is multiplied with X (and Y comes from y), so I think that I need to backprop through Y, right?
Cheers!