@ptrblck I see people using logits like this for KL divergence loss:
both pred_x and pred_x_h are logits of same dimensions, applying softmax is converting them into probablilities.
pred_x = F.softmax(model(x), dim=1)
pred_x_h = F.log_softmax(model(x_h), dim=1)
F.kl_div(pred_x_h, pred_x, None, None, reduction=‘sum’).
I am new to pytorch, not sure if thats the right thing to do?