How to calculate js Divergence for two probability distributions in PyTorch? or how to add two distribution in pytorch?