Double softmax

[bkj]

On this line you're applying a softmax to the similarities.

Then later you apply cross_entropy, which is a log softmax + NLL loss.

I think you probably want to remove the first softmax.

~ Ben