Error in normalization of losses

[alexrame]

Hey, thanks for the nice work.
Just reading your code, and I think there is a mistake in normalizing logits without the true label.
probs /= (all_probs.sum(dim=1, keepdim=True) + 1e-8)
should rather be
probs /= (probs.sum(dim=1, keepdim=True) + 1e-8)
in relation_loss in the ensembles.py file

And then similarly, for example for cos reg_type,
probs = probs / torch.sqrt(((all_probs ** 2).sum(dim=1, keepdim=True) + 1e-8)) # l2 normed
should be
probs = probs / torch.sqrt(((probs ** 2).sum(dim=1, keepdim=True) + 1e-8)) # l2 normed
as far as I understood your paper