Solved by tom in post #2 
              
              
                Less brief definitions of the Kullback-Leibler divergence spell out the that the term is to be considered to be 0 when log y_i is 0. 
This is by continuous extension: 
[image] 
The PyTorch implementation uses torch.where to achieve this. 
You cannot get away with x_i being zero when y_i is not, the…