using all singular values

[arminarj]

adding @jameschapman19 suggestions to the loss, plus zero checking for torch.sqrt(0) nan values.

Dear @Michaelvll please double-check it before merging it.

Armin

[jameschapman19]

Hi both, I've changed my mind on this. I now think that the use all singular values route is just wrong and would be better to push everything through topk (and use_all_sing_val could just default to k=max). Again referring to Andrew 2013 (eqn 10): the tracenorm of T is given by the trace of the matrix square root of T'T. This doesn't to me have a clear relation to the diagonal elements of T'T - perhaps with some conditions e.g. if T'T were diagonal but not the conditions here.

[jameschapman19]

Open to disagreement - I've got in touch with the first person to do it in this way 6 years ago which seems to be the foundation of (what I think is) the bug to see if there's a logic ive missed