Onehot was a bit buggy in some edge case (I think it was when only 1 class)
The onehot function expected to already have unsqueeze the onehot dimension, but this was confusing and let to bugs when I fixed the first issue. So, I removed the labels = labels.unsqueeze(1) step and just let the onehot handle it
unsqueezethe onehot dimension, but this was confusing and let to bugs when I fixed the first issue. So, I removed thelabels = labels.unsqueeze(1)step and just let the onehot handle it